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

@@ -35,9 +35,16 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.resumable.ChunkUploadRequest;
+import com.google.api.gax.resumable.ChunkUploadResponse;
 import com.google.api.gax.resumable.ResumableUploadClient;
 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.ScheduledRetryingExecutor;
 import java.io.InputStream;
+import java.time.Duration;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -54,9 +61,19 @@
 public class ResumableUploadCallableImpl<RequestT, ResponseT>
     extends ResumableUploadCallable<RequestT, ResponseT> {
 
+  private static final RetrySettings RETRY_SETTINGS =
+      RetrySettings.newBuilder()
+          .setInitialRetryDelayDuration(Duration.ofMillis(100))
+          .setRetryDelayMultiplier(1.3)
+          .setMaxRetryDelayDuration(Duration.ofMinutes(1))
+          .setMaxAttempts(5)
+          .build();
+
   private final ResumableUploadClient<RequestT, ResponseT> client;
   private final ResumableUploadCallSettings defaultCallSettings;
   private final ClientContext clientContext;
+  private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      retryingUploadChunkCallable;
 
   public ResumableUploadCallableImpl(
       ResumableUploadClient<RequestT, ResponseT> client,
@@ -65,7 +82,13 @@ public ResumableUploadCallableImpl(
     this.client = checkNotNull(client, "client must not be null");
     this.defaultCallSettings =
         checkNotNull(defaultCallSettings, "defaultCallSettings must not be null");
-    this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
+    this.clientContext =
+        checkNotNull(clientContext, "clientContext must not be null").toBuilder()
+            .setDefaultCallContext(
+                clientContext.getDefaultCallContext().withRetrySettings(RETRY_SETTINGS))
+            .build();
+    this.retryingUploadChunkCallable =
+        createRetryingCallable(client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD);
   }
 
   @Override
@@ -87,16 +110,24 @@ public ResumableUploadFuture<ResponseT> futureCall(
     }
 
     return ResumableUploadFutureImpl.create(
-        startFuture,
-        client.uploadChunkCallable(),
-        payload,
-        effectiveSettings,
-        clientContext.getDefaultCallContext());
+        startFuture, retryingUploadChunkCallable, payload, effectiveSettings, clientContext);
   }
 
   @Override
   public ResumableUploadFuture<ResponseT> resumeCall(
       String sessionUrl, InputStream payload, @Nullable ResumableUploadCallSettings settings) {
     throw new UnsupportedOperationException("Session resumption is not yet implemented.");
   }
+
+  private <ReqT, RespT> UnaryCallable<ReqT, RespT> createRetryingCallable(
+      UnaryCallable<ReqT, RespT> callable, ResumableUploadCommand command) {
+    RetryAlgorithm<RespT> retryAlgorithm =
+        new RetryAlgorithm<>(
+            new ResumableUploadResultRetryAlgorithm<>(command),
+            new ExponentialRetryAlgorithm(RETRY_SETTINGS, clientContext.getClock()));
+    return new RetryingCallable<>(
+        clientContext.getDefaultCallContext(),
+        checkNotNull(callable, "callable must not be null"),
+        new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor()));
+  }
 }

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

@@ -45,6 +45,7 @@
 import java.io.InputStream;
 import java.util.Arrays;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.Executor;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -59,6 +60,9 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
 
   private static final byte[] EMPTY_PAYLOAD = new byte[0];
 
+  private final Executor chunkExecutor =
+      MoreExecutors.newSequentialExecutor(MoreExecutors.directExecutor());
+
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       uploadChunkCallable;
   private final String uploadUrl;
@@ -74,13 +78,14 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
       String uploadUrl,
       InputStream payload,
       int chunkSize,
-      ApiCallContext callContext) {
+      ClientContext clientContext) {
     this.uploadChunkCallable =
         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.callContext = clientContext.getDefaultCallContext();
     this.buffer = new byte[chunkSize];
   }
 
@@ -93,7 +98,7 @@ ApiFuture<ResponseT> start() {
           }
         },
         MoreExecutors.directExecutor());
-    transmitChunk(0L);
+    chunkExecutor.execute(() -> transmitChunk(0L));
     return result;
   }
 
@@ -157,9 +162,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);
+                chunkExecutor.execute(() -> transmitChunk(nextOffset));
               }
             }
 

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

@@ -67,7 +67,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
       uploadChunkCallable;
   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;
@@ -87,10 +87,10 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
       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, clientContext);
     try {
       future.start();
     } catch (Throwable t) {
@@ -104,14 +104,14 @@ private ResumableUploadFutureImpl(
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
-      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");
     this.inFlightFuture = startFuture;
   }
 
@@ -131,7 +131,7 @@ public void onSuccess(ResumableUploadSession session) {
                     uploadSessionUrl,
                     payload,
                     settings.getChunkSize(),
-                    callContext);
+                    clientContext);
             ApiFuture<ResponseT> uploadFuture;
             try {
               uploadFuture = coordinator.start();

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java

@@ -214,6 +214,8 @@ void testWatched_usesJavaTimeMethods() {
   void testResumableUploadCallable() {
     ResumableUploadClient<String, String> uploadClient =
         mock(ResumableUploadClient.class, Mockito.withSettings().withoutAnnotations());
+    when(uploadClient.uploadChunkCallable())
+        .thenReturn(mock(UnaryCallable.class, Mockito.withSettings().withoutAnnotations()));
     ResumableUploadCallSettings settings =
         ResumableUploadCallSettings.newBuilder().setChunkSize(1024).build();
 

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java

@@ -404,6 +404,102 @@ void testResumeCall_throwsUnsupportedOperationException() {
         () -> callable.resumeCall("https://upload.url/session", streamOf("data"), null));
   }
 
+  @Test
+  void testChunkRetry_transientFailureThenSuccess_retriesAndSucceeds() throws Exception {
+    stubStartSession("https://upload.url/chunk-retry-ok");
+    TrackableStream stream = new TrackableStream("01234567"); // exactly 1 chunk of 8 bytes
+    when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any()))
+        .thenReturn(
+            ApiFutures.immediateFailedFuture(createApiException(503, StatusCode.Code.UNAVAILABLE)))
+        .thenReturn(
+            ApiFutures.immediateFuture(
+                ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "chunk-done")));
+
+    ResumableUploadFuture<String> future = callable.futureCall("resource-path", stream, null);
+
+    assertThat(future.get()).isEqualTo("chunk-done");
+    assertThat(future.isDone()).isTrue();
+    assertThat(stream.totalBytesRead).isEqualTo(8);
+    assertThat(stream.closed).isTrue();
+
+    ArgumentCaptor<ChunkUploadRequest> captor = ArgumentCaptor.forClass(ChunkUploadRequest.class);
+    verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any());
+    List<ChunkUploadRequest> requests = captor.getAllValues();
+    assertThat(requests.get(0).getOffset()).isEqualTo(0);
+    assertThat(requests.get(0).getPayload()).isEqualTo("01234567".getBytes(StandardCharsets.UTF_8));
+    assertThat(requests.get(1).getOffset()).isEqualTo(0);
+    assertThat(requests.get(1).getPayload()).isEqualTo("01234567".getBytes(StandardCharsets.UTF_8));
+  }
+
+  @Test
+  void testChunkRetry_transientFailureExhaustion_surfacesLastError() {
+    stubStartSession("https://upload.url/chunk-exhaustion");
+    when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any()))
+        .thenReturn(
+            ApiFutures.immediateFailedFuture(createApiException(503, StatusCode.Code.UNAVAILABLE)));
+
+    ResumableUploadFuture<String> future =
+        callable.futureCall("resource-path", streamOf("hello"), null);
+
+    ExecutionException exception = assertThrows(ExecutionException.class, future::get);
+    assertThat(exception.getCause()).isInstanceOf(ApiException.class);
+    assertThat(((ApiException) exception.getCause()).getStatusCode().getTransportCode())
+        .isEqualTo(503);
+
+    // Default chunk retry settings has maxAttempts = 5
+    verify(mockChunkCallable, times(5)).futureCall(any(), any());
+  }
+
+  @Test
+  void testChunkRetry_cancellationDuringBackoff_deschedulesPendingAttempt() {
+    stubStartSession("https://upload.url/cancel-backoff");
+    SettableApiFuture<ChunkUploadResponse<String>> chunkAttempt0Future = SettableApiFuture.create();
+    when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any()))
+        .thenReturn(chunkAttempt0Future)
+        .thenReturn(
+            ApiFutures.immediateFuture(
+                ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "should-not-reach")));
+
+    ResumableUploadFuture<String> sessionFuture =
+        callable.futureCall("resource-path", streamOf("hello"), null);
+
+    // Fail attempt 0 with 503 to schedule backoff
+    chunkAttempt0Future.setException(createApiException(503, StatusCode.Code.UNAVAILABLE));
+
+    // Cancel while backoff is pending
+    assertThat(sessionFuture.cancel(true)).isTrue();
+    assertThat(sessionFuture.isCancelled()).isTrue();
+    assertThrows(CancellationException.class, sessionFuture::get);
+
+    // Only attempt 0 occurred; attempt 1 was de-scheduled
+    verify(mockChunkCallable, times(1)).futureCall(any(), any());
+  }
+
+  private static class HttpStatusStatusCode implements StatusCode {
+    private final int httpStatus;
+    private final StatusCode.Code code;
+
+    HttpStatusStatusCode(int httpStatus, StatusCode.Code code) {
+      this.httpStatus = httpStatus;
+      this.code = code;
+    }
+
+    @Override
+    public StatusCode.Code getCode() {
+      return code;
+    }
+
+    @Override
+    public Integer getTransportCode() {
+      return httpStatus;
+    }
+  }
+
+  private static ApiException createApiException(int httpStatus, StatusCode.Code code) {
+    return ApiExceptionFactory.createException(
+        "HTTP " + httpStatus, null, new HttpStatusStatusCode(httpStatus, code), false);
+  }
+
   private void stubStartSession(String uploadUrl) {
     when(mockStartCallable.futureCall(any(), any()))
         .thenReturn(
@@ -424,11 +520,21 @@ private static void assertChunk(
 
   private static class TrackableStream extends ByteArrayInputStream {
     boolean closed = false;
+    int totalBytesRead = 0;
 
     TrackableStream(String content) {
       super(content.getBytes(StandardCharsets.UTF_8));
     }
 
+    @Override
+    public int read(byte[] b, int off, int len) {
+      int read = super.read(b, off, len);
+      if (read > 0) {
+        totalBytesRead += read;
+      }
+      return read;
+    }
+
     @Override
     public void close() throws IOException {
       closed = true;