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

@@ -39,13 +39,21 @@
 import com.google.api.gax.resumable.ChunkUploadRequest;
 import com.google.api.gax.resumable.ChunkUploadResponse;
 import com.google.api.gax.resumable.ResumableUploadSession;
+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 com.google.errorprone.annotations.concurrent.GuardedBy;
 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;
 
@@ -57,18 +65,34 @@
 @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 Object lock = new Object();
+  private final AtomicBoolean dispatching = new AtomicBoolean(false);
+  private final AtomicLong nextChunkOffset = new AtomicLong(-1L);
 
   private final SettableApiFuture<ResponseT> result;
   private final ApiFuture<ResumableUploadSession> startFuture;
-  private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
-      uploadChunkCallable;
+  private final RetryingCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      retryingChunkCallable;
   private final InputStream payload;
   private final byte[] buffer;
   private final int chunkSize;
   private final ApiCallContext callContext;
+  private final ClientContext clientContext;
+  private final RetrySettings chunkRetrySettings;
 
   private volatile @Nullable String uploadSessionUrl;
 
@@ -87,17 +111,30 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
-      ApiCallContext callContext) {
+      ApiCallContext callContext,
+      ClientContext clientContext) {
     this.result = checkNotNull(result, "result must not be null");
     this.startFuture = checkNotNull(startFuture, "startFuture must not be null");
-    this.uploadChunkCallable =
-        checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
+    checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.payload = checkNotNull(payload, "payload must not be null");
     checkNotNull(settings, "settings must not be null");
     checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0");
     this.chunkSize = settings.getChunkSize();
     this.callContext = checkNotNull(callContext, "callContext must not be null");
+    this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
+    this.chunkRetrySettings = DEFAULT_CHUNK_RETRY_SETTINGS;
     this.buffer = new byte[chunkSize];
+
+    RetryAlgorithm<ChunkUploadResponse<ResponseT>> retryAlgorithm =
+        new RetryAlgorithm<>(
+            new UploadResultRetryAlgorithm<>(UploadCommand.UPLOAD),
+            new ExponentialRetryAlgorithm(chunkRetrySettings, clientContext.getClock()));
+    ScheduledRetryingExecutor<ChunkUploadResponse<ResponseT>> retryingExecutor =
+        new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
+    this.retryingChunkCallable =
+        new RetryingCallable<>(
+            clientContext.getDefaultCallContext(), uploadChunkCallable, retryingExecutor);
+
     synchronized (lock) {
       this.inFlightFuture = startFuture;
     }
@@ -115,7 +152,7 @@ public void onSuccess(ResumableUploadSession session) {
               }
             }
             uploadSessionUrl = session.getUploadUrl();
-            transmitChunk(0L);
+            scheduleNextChunk(0L);
           }
 
           @Override
@@ -197,7 +234,24 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) {
     }
   }
 
-  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) {
     synchronized (lock) {
       if (done) {
         return;
@@ -236,47 +290,44 @@ private void transmitChunk(long currentOffset) {
             .setFinal(isFinal)
             .build();
 
-    long chunkLength = chunkPayload.length;
-    try {
-      ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
-          uploadChunkCallable.futureCall(chunkRequest, callContext);
-      setInFlightFuture(chunkFuture);
+    RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture =
+        retryingChunkCallable.futureCall(chunkRequest, callContext);
+    setInFlightFuture(retryingFuture);
 
-      ApiFutures.addCallback(
-          chunkFuture,
-          new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
-            @Override
-            public void onSuccess(ChunkUploadResponse<ResponseT> response) {
-              synchronized (lock) {
-                if (done) {
-                  return;
-                }
-              }
-              long nextOffset = currentOffset + chunkLength;
-              if (response.isComplete()) {
-                finish(response.getResponse(), null);
-              } else if (isFinal) {
-                finish(
-                    null,
-                    new IllegalStateException(
-                        "Upload stream ended and final chunk was transmitted, but server returned"
-                            + " incomplete status"));
-              } else {
-                transmitChunk(nextOffset);
+    long chunkLength = chunkPayload.length;
+    ApiFutures.addCallback(
+        retryingFuture,
+        new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
+          @Override
+          public void onSuccess(ChunkUploadResponse<ResponseT> response) {
+            synchronized (lock) {
+              if (done) {
+                return;
               }
             }
+            long nextOffset = currentOffset + chunkLength;
+            if (response.isComplete()) {
+              finish(response.getResponse(), null);
+            } else if (isFinal) {
+              finish(
+                  null,
+                  new IllegalStateException(
+                      "Upload stream ended and final chunk was transmitted, but server returned"
+                          + " incomplete status for upload URL: "
+                          + url));
+            } else {
+              scheduleNextChunk(nextOffset);
+            }
+          }
 
-            @Override
-            public void onFailure(Throwable t) {
-              if (t instanceof CancellationException) {
-                return;
-              }
-              finish(null, t);
+          @Override
+          public void onFailure(Throwable t) {
+            if (t instanceof CancellationException) {
+              return;
             }
-          },
-          MoreExecutors.directExecutor());
-    } catch (Throwable t) {
-      finish(null, t);
-    }
+            finish(null, t);
+          }
+        },
+        MoreExecutors.directExecutor());
   }
 }

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

@@ -61,10 +61,32 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
       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) {
     SettableApiFuture<ResponseT> result = SettableApiFuture.create();
     ResumableUploadChunkCoordinator<ResponseT> coordinator =
         new ResumableUploadChunkCoordinator<>(
-            result, startFuture, uploadChunkCallable, payload, settings, callContext);
+            result,
+            startFuture,
+            uploadChunkCallable,
+            payload,
+            settings,
+            callContext,
+            clientContext);
     ResumableUploadFutureImpl<ResponseT> handle =
         new ResumableUploadFutureImpl<>(result, coordinator);
     coordinator.start();