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

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));