feat(gax): retry chunk upload on transient errors

Wraps chunk uploads in a retrying executor to retry transient network and server errors using exponential backoff. Retries individual chunks without restarting the entire upload session.

Author: whowes

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java

@@ -91,7 +91,8 @@ public ResumableUploadFuture<ResponseT> futureCall(
         client.uploadChunkCallable(),
         payload,
         effectiveSettings,
-        clientContext.getDefaultCallContext());
+        clientContext.getDefaultCallContext(),
+        clientContext);
   }
 
   @Override

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java

@@ -38,12 +38,20 @@
 import com.google.api.gax.resumable.ChunkUploadRequest;
 import com.google.api.gax.resumable.ChunkUploadResponse;
 import com.google.api.gax.resumable.ResumableUploadStatus;
+import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
+import com.google.api.gax.retrying.RetryAlgorithm;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.retrying.RetryingFuture;
+import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import com.google.common.io.ByteStreams;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.io.IOException;
 import java.io.InputStream;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -55,10 +63,25 @@
 @NullMarked
 final class ResumableUploadChunkCoordinator<ResponseT> {
 
+  static final RetrySettings DEFAULT_CHUNK_RETRY_SETTINGS =
+      RetrySettings.newBuilder()
+          .setInitialRetryDelayDuration(Duration.ofMillis(100))
+          .setRetryDelayMultiplier(1.3)
+          .setMaxRetryDelayDuration(Duration.ofMinutes(1))
+          .setInitialRpcTimeoutDuration(Duration.ofSeconds(30))
+          .setRpcTimeoutMultiplier(1.0)
+          .setMaxRpcTimeoutDuration(Duration.ofSeconds(30))
+          .setTotalTimeoutDuration(Duration.ofMinutes(5))
+          .setMaxAttempts(5)
+          .build();
+
   private static final byte[] EMPTY_PAYLOAD = new byte[0];
 
-  private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
-      uploadChunkCallable;
+  private final AtomicBoolean dispatching = new AtomicBoolean(false);
+  private final AtomicLong nextChunkOffset = new AtomicLong(-1L);
+
+  private final RetryingCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      retryingChunkCallable;
   private final String uploadUrl;
   private final InputStream payload;
   private final byte[] buffer;
@@ -72,14 +95,25 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
       String uploadUrl,
       InputStream payload,
       int chunkSize,
-      ApiCallContext callContext) {
-    this.uploadChunkCallable =
-        checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
+      ApiCallContext callContext,
+      ClientContext clientContext) {
+    checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
     this.payload = checkNotNull(payload, "payload must not be null");
     this.chunkSize = chunkSize;
     this.callContext = checkNotNull(callContext, "callContext must not be null");
+    checkNotNull(clientContext, "clientContext must not be null");
     this.buffer = new byte[chunkSize];
+
+    RetryAlgorithm<ChunkUploadResponse<ResponseT>> retryAlgorithm =
+        new RetryAlgorithm<>(
+            new ResumableUploadResultRetryAlgorithm<>(ResumableUploadCommand.UPLOAD),
+            new ExponentialRetryAlgorithm(DEFAULT_CHUNK_RETRY_SETTINGS, clientContext.getClock()));
+    ScheduledRetryingExecutor<ChunkUploadResponse<ResponseT>> retryingExecutor =
+        new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
+    this.retryingChunkCallable =
+        new RetryingCallable<>(
+            clientContext.getDefaultCallContext(), uploadChunkCallable, retryingExecutor);
   }
 
   ApiFuture<ResponseT> start() {
@@ -93,11 +127,28 @@ ApiFuture<ResponseT> start() {
           }
         },
         MoreExecutors.directExecutor());
-    transmitChunk(0L);
+    scheduleNextChunk(0L);
     return result;
   }
 
-  private void transmitChunk(long currentOffset) {
+  private void scheduleNextChunk(long offset) {
+    nextChunkOffset.set(offset);
+    if (dispatching.compareAndSet(false, true)) {
+      driveLoop();
+    }
+  }
+
+  private void driveLoop() {
+    do {
+      long offset = nextChunkOffset.getAndSet(-1L);
+      if (offset >= 0) {
+        transmitSingleChunk(offset);
+      }
+      dispatching.set(false);
+    } while (nextChunkOffset.get() >= 0 && dispatching.compareAndSet(false, true));
+  }
+
+  private void transmitSingleChunk(long currentOffset) {
     if (result.isDone()) {
       return;
     }
@@ -130,12 +181,12 @@ private void transmitChunk(long currentOffset) {
 
     long chunkLength = chunkPayload.length;
     try {
-      ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
-          uploadChunkCallable.futureCall(chunkRequest, callContext);
-      this.currentChunkFuture = chunkFuture;
+      RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture =
+          retryingChunkCallable.futureCall(chunkRequest, callContext);
+      this.currentChunkFuture = retryingFuture;
 
       ApiFutures.addCallback(
-          chunkFuture,
+          retryingFuture,
           new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
             @Override
             public void onSuccess(ChunkUploadResponse<ResponseT> response) {
@@ -149,9 +200,10 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
                 result.setException(
                     new IllegalStateException(
                         "Upload stream ended and final chunk was transmitted, but server returned"
-                            + " incomplete status"));
+                            + " incomplete status for upload URL: "
+                            + uploadUrl));
               } else {
-                transmitChunk(nextOffset);
+                scheduleNextChunk(nextOffset);
               }
             }
 

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java

@@ -68,6 +68,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   private final InputStream payload;
   private final ResumableUploadCallSettings settings;
   private final ApiCallContext callContext;
+  private final ClientContext clientContext;
   private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
 
   private volatile @Nullable String uploadSessionUrl;
@@ -81,22 +82,31 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   @GuardedBy("lock")
   private @Nullable ApiFuture<?> inFlightFuture;
 
-  /**
-   * Creates and initiates a new resumable upload future tracking session initiation and chunk
-   * streaming.
-   *
-   * <p>The provided {@code payload} stream is managed by the returned future and will be closed
-   * automatically upon completion, failure, or cancellation.
-   */
   static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
       ApiFuture<ResumableUploadSession> startFuture,
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
       ApiCallContext callContext) {
+    return create(
+        startFuture,
+        uploadChunkCallable,
+        payload,
+        settings,
+        callContext,
+        ClientContext.newBuilder().setDefaultCallContext(callContext).build());
+  }
+
+  static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
+      ApiFuture<ResumableUploadSession> startFuture,
+      UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      InputStream payload,
+      ResumableUploadCallSettings settings,
+      ApiCallContext callContext,
+      ClientContext clientContext) {
     ResumableUploadFutureImpl<ResponseT> future =
         new ResumableUploadFutureImpl<>(
-            startFuture, uploadChunkCallable, payload, settings, callContext);
+            startFuture, uploadChunkCallable, payload, settings, callContext, clientContext);
     future.start();
     return future;
   }
@@ -106,14 +116,16 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
-      ApiCallContext callContext) {
+      ApiCallContext callContext,
+      ClientContext clientContext) {
     this.startFuture = checkNotNull(startFuture, "startFuture must not be null");
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.payload = checkNotNull(payload, "payload must not be null");
     this.settings = checkNotNull(settings, "settings must not be null");
     checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0");
     this.callContext = checkNotNull(callContext, "callContext must not be null");
+    this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
     synchronized (lock) {
       this.inFlightFuture = startFuture;
     }
@@ -137,7 +149,8 @@ public void onSuccess(ResumableUploadSession session) {
                       uploadSessionUrl,
                       payload,
                       settings.getChunkSize(),
-                      callContext);
+                      callContext,
+                      clientContext);
               uploadFuture = coordinator.start();
               inFlightFuture = uploadFuture;
             }