fix(spanner): honor maxAttempts and totalTimeout in streaming resume loop (#14370)

The resume loop in ResumableStreamIterator restarted a broken stream
indefinitely for any retryable error, ignoring the maxAttempts and
totalTimeout configured in the retry settings for ExecuteStreamingSql.

A streaming query could therefore retry forever when the server kept
returning a retryable error, for example when a user configured
DEADLINE_EXCEEDED as a retryable code and every attempt timed out.

The loop now counts consecutive failed attempts and stops retrying,
rethrowing the last exception, when the configured maxAttempts is
reached. It also enforces the configured totalTimeout as a wall-clock
budget for a sequence of consecutive failed attempts, measured from the
first failure of the sequence: a retry is only allowed when the
retry delay still fits in the remaining budget. This applies both to
delays from the exponential backoff and to server-supplied retry delays
(RetryInfo), which previously bypassed the backoff completely.

A totalTimeout of zero means that no time budget has been set, in which
case only maxAttempts limits the retries, mirroring GAX.

Both limits only bind for custom retry settings: the default streaming
retry settings do not set maxAttempts and keep the existing unbounded
resume behavior. Progress on the stream resets both budgets, where
progress means receiving a resume token that differs from the last seen
token, so long-running streams that regularly make progress are
not terminated by an occasional transient error, while a stream that
keeps returning the same token cannot reset the budget indefinitely.

Author: rahul2393

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/ResumableStreamIterator.java

@@ -28,7 +28,6 @@
 import com.google.api.gax.retrying.RetrySettings;
 import com.google.api.gax.rpc.StatusCode.Code;
 import com.google.cloud.spanner.AbstractResultSet.CloseableIterator;
-import com.google.cloud.spanner.v1.stub.SpannerStubSettings;
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.AbstractIterator;
@@ -38,7 +37,6 @@
 import io.opentelemetry.api.common.Attributes;
 import java.io.IOException;
 import java.util.LinkedList;
-import java.util.Objects;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
@@ -52,12 +50,28 @@
  * track of the most recent resume token seen, and will buffer partial result set chunks that do not
  * have a resume token until one is seen or buffer space is exceeded, which reduces the chance of
  * yielding data to the caller that cannot be resumed.
+ *
+ * <p>The resume loop is bounded by the streaming retry settings: consecutive failed attempts are
+ * limited by {@link RetrySettings#getMaxAttempts()} if that has been set to a value greater than
+ * zero, and the total time spent on a sequence of consecutive failed attempts is limited by {@link
+ * RetrySettings#getTotalTimeout()} if that has been set to a positive value. Only consecutive
+ * failures count against these budgets: any progress on the stream (that is, receiving a resume
+ * token that differs from the last seen resume token) resets both, so a long-running stream that
+ * regularly makes progress is not terminated by an occasional transient error. The default settings
+ * have no maximum number of attempts or total timeout. Setting maxAttempts to 1 disables streaming
+ * retries.
+ *
+ * <p>These limits bound the number of streams that this iterator starts. Each (re)started stream is
+ * a call through the underlying GAX callable, and GAX applies the same {@link RetrySettings} to
+ * attempts of that call that fail before any response has been received. A stream that repeatedly
+ * fails before its first response can therefore consist of up to maxAttempts RPC attempts itself,
+ * so a configured maxAttempts of N bounds the total number of RPC attempts without progress by N*N,
+ * not by N. The total timeout is measured in wall-clock time from the first failure of the sequence
+ * and therefore spans both layers.
  */
 @VisibleForTesting
 abstract class ResumableStreamIterator extends AbstractIterator<PartialResultSet>
     implements CloseableIterator<PartialResultSet> {
-  private static final RetrySettings DEFAULT_STREAMING_RETRY_SETTINGS =
-      SpannerStubSettings.newBuilder().executeStreamingSqlSettings().getRetrySettings();
   private final ErrorHandler errorHandler;
   private AsyncResultSet.StreamMessageListener streamMessageListener;
   private final RetrySettings streamingRetrySettings;
@@ -69,7 +83,20 @@ abstract class ResumableStreamIterator extends AbstractIterator<PartialResultSet
   private final ISpan span;
   private final TraceWrapper tracer;
   private CloseableIterator<PartialResultSet> stream;
+
+  /**
+   * The number of consecutive failed attempts without any progress on the stream. Reset to zero
+   * every time the stream returns a new resume token.
+   */
   private int attempts;
+
+  /**
+   * The value of {@link System#nanoTime()} at the first failure of the current sequence of
+   * consecutive failed attempts. Only meaningful when {@link #attempts} is nonzero. Used to enforce
+   * {@link RetrySettings#getTotalTimeout()} when a positive timeout is configured.
+   */
+  private long retrySequenceStartNanos = -1L;
+
   private ByteString resumeToken;
   private boolean finished;
   private final XGoogSpannerRequestId requestId;
@@ -123,17 +150,42 @@ protected ResumableStreamIterator(
     this.requestId = xGoogRequestIdCreator.nextRequestId(0);
   }
 
-  private ExponentialBackOff newBackOff() {
-    if (Objects.equals(streamingRetrySettings, DEFAULT_STREAMING_RETRY_SETTINGS)) {
-      return new ExponentialBackOff.Builder()
-          .setMultiplier(streamingRetrySettings.getRetryDelayMultiplier())
-          .setInitialIntervalMillis(
-              Math.max(10, (int) streamingRetrySettings.getInitialRetryDelay().toMillis()))
-          .setMaxIntervalMillis(
-              Math.max(1000, (int) streamingRetrySettings.getMaxRetryDelay().toMillis()))
-          .setMaxElapsedTimeMillis(Integer.MAX_VALUE) // Prevent Backoff.STOP from getting returned.
-          .build();
+  /**
+   * Returns true if the number of consecutive failed attempts has reached the maximum number of
+   * attempts in the retry settings. {@link RetrySettings#getMaxAttempts()} equal to zero means that
+   * no maximum has been set, and that the number of attempts is unlimited. This is also the value
+   * in the default streaming retry settings, which means that only users who have explicitly opted
+   * in to a maximum number of attempts are affected by this limit.
+   */
+  private boolean maxAttemptsExhausted() {
+    int maxAttempts = streamingRetrySettings.getMaxAttempts();
+    return maxAttempts > 0 && attempts >= maxAttempts;
+  }
+
+  /**
+   * Returns true if retrying after the proposed delay would exceed the total timeout in the retry
+   * settings. The total timeout limits the wall-clock time that is spent on a sequence of
+   * consecutive failed attempts without progress, measured from the first failure of the sequence.
+   * It is only enforced for retry settings that set a positive total timeout: a total timeout of
+   * zero means that no total timeout has been set, and that only maxAttempts (if set) limits the
+   * retries. This mirrors the interpretation of these values in GAX.
+   */
+  private boolean totalTimeoutExceeded(long proposedDelayMillis) {
+    long totalTimeoutMillis = streamingRetrySettings.getTotalTimeout().toMillis();
+    if (totalTimeoutMillis <= 0L) {
+      return false;
+    }
+    long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(nanoTime() - retrySequenceStartNanos);
+    if (elapsedMillis < 0L) {
+      elapsedMillis = 0L;
+    }
+    if (elapsedMillis >= totalTimeoutMillis) {
+      return true;
     }
+    return Math.max(proposedDelayMillis, 0L) >= totalTimeoutMillis - elapsedMillis;
+  }
+
+  private ExponentialBackOff newBackOff() {
     return new ExponentialBackOff.Builder()
         .setMultiplier(streamingRetrySettings.getRetryDelayMultiplier())
         // All of these values must be > 0.
@@ -150,19 +202,14 @@ private ExponentialBackOff newBackOff() {
                 (int)
                     Math.min(
                         streamingRetrySettings.getMaxRetryDelay().toMillis(), Integer.MAX_VALUE)))
-        .setMaxElapsedTimeMillis(
-            Math.max(
-                1,
-                (int)
-                    Math.min(
-                        streamingRetrySettings.getTotalTimeout().toMillis(), Integer.MAX_VALUE)))
+        // The total timeout is enforced explicitly in computeNext(), based on the elapsed time
+        // since the first failure of the current retry sequence. Prevent the backoff from
+        // returning BackOff.STOP, as that would misinterpret a total timeout of zero (that is, no
+        // total timeout) as a total timeout of one millisecond.
+        .setMaxElapsedTimeMillis(Integer.MAX_VALUE)
         .build();
   }
 
-  private void backoffSleep(Context context, BackOff backoff) throws SpannerException {
-    backoffSleep(context, nextBackOffMillis(backoff));
-  }
-
   private static long nextBackOffMillis(BackOff backoff) throws SpannerException {
     try {
       return backoff.nextBackOffMillis();
@@ -263,6 +310,14 @@ protected PartialResultSet computeNext() {
           PartialResultSet next = stream.next();
           boolean hasResumeToken = !next.getResumeToken().isEmpty();
           if (hasResumeToken) {
+            // Only a resume token that differs from the last seen token is progress: a stream
+            // that repeatedly returns the token that was used to resume it has not moved past it.
+            if (!next.getResumeToken().equals(resumeToken)) {
+              // The stream made progress, so reset the budget for consecutive failed attempts.
+              attempts = 0;
+              backOff = null;
+              retrySequenceStartNanos = -1L;
+            }
             resumeToken = next.getResumeToken();
             safeToRetry = true;
           }
@@ -287,25 +342,7 @@ protected PartialResultSet computeNext() {
         }
       } catch (SpannerException spannerException) {
         if (safeToRetry && isRetryable(spannerException)) {
-          span.addAnnotation("Stream broken. Safe to retry", spannerException);
-          logger.log(Level.FINE, "Retryable exception, will sleep and retry", spannerException);
-          // Truncate any items in the buffer before the last retry token.
-          while (!buffer.isEmpty() && buffer.getLast().getResumeToken().isEmpty()) {
-            buffer.removeLast();
-          }
-          assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken);
-          stream = null;
-          try (IScope s = tracer.withSpan(span)) {
-            long delay = spannerException.getRetryDelayInMillis();
-            if (delay != -1) {
-              backoffSleep(context, delay);
-            } else {
-              if (this.backOff == null) {
-                this.backOff = newBackOff();
-              }
-              backoffSleep(context, this.backOff);
-            }
-          }
+          handleRetryableException(context, spannerException);
 
           continue;
         }
@@ -331,6 +368,62 @@ && prepareIteratorForRetryOnDifferentGrpcChannel()) {
     }
   }
 
+  /** Monotonic time source, overridable for deterministic retry-budget tests. */
+  @VisibleForTesting
+  long nanoTime() {
+    return System.nanoTime();
+  }
+
+  @VisibleForTesting
+  long checkRetryBudgetAndGetDelay(SpannerException spannerException) {
+    if (attempts == 0) {
+      retrySequenceStartNanos = nanoTime();
+    }
+    attempts++;
+    if (maxAttemptsExhausted()) {
+      span.addAnnotation(
+          "Stream broken. Not retrying because the maximum number of attempts has been"
+              + " exhausted",
+          spannerException);
+      span.setStatus(spannerException);
+      throw spannerException;
+    }
+    // Determine the retry delay: either the delay that the server included in the error, or
+    // otherwise a delay determined by the exponential backoff.
+    long delayMillis = spannerException.getRetryDelayInMillis();
+    if (delayMillis == -1L) {
+      if (this.backOff == null) {
+        this.backOff = newBackOff();
+      }
+      delayMillis = nextBackOffMillis(this.backOff);
+    }
+    // The total timeout budget applies regardless of whether the delay came from the server
+    // or from the backoff.
+    if (totalTimeoutExceeded(delayMillis)) {
+      span.addAnnotation(
+          "Stream broken. Not retrying because the total timeout has been exhausted",
+          spannerException);
+      span.setStatus(spannerException);
+      throw spannerException;
+    }
+    return delayMillis;
+  }
+
+  private void handleRetryableException(Context context, SpannerException spannerException) {
+    long delayMillis = checkRetryBudgetAndGetDelay(spannerException);
+    span.addAnnotation("Stream broken. Safe to retry", spannerException);
+    logger.log(Level.FINE, "Retryable exception, will sleep and retry", spannerException);
+    // Truncate any items in the buffer before the last retry token.
+    while (!buffer.isEmpty() && buffer.getLast().getResumeToken().isEmpty()) {
+      buffer.removeLast();
+    }
+    assert buffer.isEmpty() || buffer.getLast().getResumeToken().equals(resumeToken);
+    stream = null;
+    try (IScope s = tracer.withSpan(span)) {
+      backoffSleep(context, delayMillis);
+    }
+  }
+
   private void startGrpcStreaming() {
     if (stream == null) {
       span.addAnnotation(

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/SpannerOptions.java

@@ -1669,6 +1669,15 @@ public Builder setRetrySettings(RetrySettings retrySettings) {
      * href="https://github.com/googleapis/googleapis/blob/master/google/spanner/v1/spanner_gapic.yaml">spanner_gapic.yaml</a>.
      * Retries are configured for idempotent methods but not for non-idempotent methods.
      *
+     * <p>For streaming queries and reads, configure {@code executeStreamingSqlSettings()} and
+     * {@code streamingReadSettings()}, respectively. Set {@code maxAttempts=1} to disable streaming
+     * retries; an empty set of retryable codes does not disable retries for intrinsically retryable
+     * errors. Defaults allow unlimited streaming resumes. When customizing retry settings, set the
+     * total timeout explicitly: calling {@code toBuilder()} on the stub's retry settings copies
+     * GAPIC's generated one-hour {@code totalTimeout}; set it to zero explicitly for unlimited
+     * resumes. Limits apply to consecutive failures and reset when the stream makes progress by
+     * returning a new resume token.
+     *
      * <p>You can set the same {@link RetrySettings} for all unary methods by calling this:
      *
      * <pre><code>

java-spanner/google-cloud-spanner/src/main/java/com/google/cloud/spanner/spi/v1/GapicSpannerRpc.java

@@ -271,6 +271,15 @@ public class GapicSpannerRpc implements SpannerRpc {
   private static final CallOptions.Key<Boolean> BASE_CONTEXT_MARKER_KEY =
       CallOptions.Key.create("BASE_CONTEXT_MARKER_KEY");
 
+  // Normalize the generated placeholder to the historical streaming resume policy.
+  static final RetrySettings DEFAULT_STREAMING_RETRY_SETTINGS =
+      RetrySettings.newBuilder()
+          .setTotalTimeoutDuration(Duration.ZERO)
+          .setMaxAttempts(0)
+          .setInitialRetryDelayDuration(Duration.ofMillis(10))
+          .setMaxRetryDelayDuration(Duration.ofMillis(1000))
+          .build();
+
   private final RequestIdCreator requestIdCreator = new RequestIdCreatorImpl();
   private boolean rpcIsClosed;
   private final SpannerStub spannerStub;
@@ -452,8 +461,14 @@ public GapicSpannerRpc(final SpannerOptions options) {
         DIRECTPATH_CHANNEL_CREATED =
             ((GrpcTransportChannel) clientContext.getTransportChannel()).isDirectPath()
                 && isEnableDirectAccess;
-        this.readRetrySettings =
+        SpannerStubSettings.Builder defaultStubSettings = SpannerStubSettings.newBuilder();
+        RetrySettings configuredReadRetrySettings =
             options.getSpannerStubSettings().streamingReadSettings().getRetrySettings();
+        this.readRetrySettings =
+            configuredReadRetrySettings.equals(
+                    defaultStubSettings.streamingReadSettings().getRetrySettings())
+                ? DEFAULT_STREAMING_RETRY_SETTINGS
+                : configuredReadRetrySettings;
         Set<Code> streamingReadRetryableCodes =
             options.getSpannerStubSettings().streamingReadSettings().getRetryableCodes();
         this.readRetryableCodes =
@@ -463,8 +478,13 @@ public GapicSpannerRpc(final SpannerOptions options) {
                     .add(Code.RESOURCE_EXHAUSTED)
                     .build()
                 : streamingReadRetryableCodes;
-        this.executeQueryRetrySettings =
+        RetrySettings configuredQueryRetrySettings =
             options.getSpannerStubSettings().executeStreamingSqlSettings().getRetrySettings();
+        this.executeQueryRetrySettings =
+            configuredQueryRetrySettings.equals(
+                    defaultStubSettings.executeStreamingSqlSettings().getRetrySettings())
+                ? DEFAULT_STREAMING_RETRY_SETTINGS
+                : configuredQueryRetrySettings;
         Set<Code> executeStreamingSqlRetryableCodes =
             options.getSpannerStubSettings().executeStreamingSqlSettings().getRetryableCodes();
         this.executeQueryRetryableCodes =

java-spanner/google-cloud-spanner/src/test/java/com/google/cloud/spanner/MockSpannerServiceImpl.java

@@ -517,6 +517,19 @@ public static SimulatedExecutionTime ofStreamException(Exception exception, long
           0, 0, Collections.singletonList(exception), false, Collections.singleton(streamIndex));
     }
 
+    /**
+     * Creates a {@link SimulatedExecutionTime} that throws the given exceptions at the given
+     * indices in the returned stream. The exceptions and stream indices are matched by position:
+     * the first exception is thrown when the first call reaches the first stream index, the second
+     * exception when the next call reaches the second stream index, and so on. The stream index is
+     * reset for each (retried) call.
+     */
+    public static SimulatedExecutionTime ofStreamExceptions(
+        Collection<? extends Exception> exceptions, Collection<Long> streamIndices) {
+      Preconditions.checkArgument(exceptions.size() == streamIndices.size());
+      return new SimulatedExecutionTime(0, 0, exceptions, false, streamIndices);
+    }
+
     public static SimulatedExecutionTime stickyDatabaseNotFoundException(String name) {
       return ofStickyException(
           SpannerExceptionFactoryTest.newStatusDatabaseNotFoundException(name));