feat(gax): retry chunk upload on transient errors (#14422)

Wraps chunk uploads in a `RetryingCallable` to retry transient errors.

Note: the retry settings used here are intentionally different from
those that will be used for starting the upload; our requirement for
this milestone is that "the retry policy [specified by the user] is only
applicable to the initial request."

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,23 @@
 public class ResumableUploadCallableImpl<RequestT, ResponseT>
     extends ResumableUploadCallable<RequestT, ResponseT> {
 
+  private static final RetrySettings RETRY_SETTINGS =
+      RetrySettings.newBuilder()
+          .setInitialRetryDelayDuration(Duration.ofSeconds(1))
+          .setRetryDelayMultiplier(2.0)
+          .setMaxRetryDelayDuration(Duration.ofSeconds(60))
+          .setMaxAttempts(5)
+          .setInitialRpcTimeoutDuration(Duration.ofMinutes(1))
+          .setRpcTimeoutMultiplier(1.0)
+          .setMaxRpcTimeoutDuration(Duration.ofMinutes(1))
+          .setTotalTimeoutDuration(Duration.ofMinutes(2))
+          .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,
@@ -66,6 +87,9 @@ public ResumableUploadCallableImpl(
     this.defaultCallSettings =
         checkNotNull(defaultCallSettings, "defaultCallSettings must not be null");
     this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
+    this.retryingUploadChunkCallable =
+        createRetryingCallable(
+            client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD, clientContext);
   }
 
   @Override
@@ -88,7 +112,7 @@ public ResumableUploadFuture<ResponseT> futureCall(
 
     return ResumableUploadFutureImpl.create(
         startFuture,
-        client.uploadChunkCallable(),
+        retryingUploadChunkCallable,
         payload,
         effectiveSettings,
         clientContext.getDefaultCallContext());
@@ -99,4 +123,18 @@ public ResumableUploadFuture<ResponseT> resumeCall(
       String sessionUrl, InputStream payload, @Nullable ResumableUploadCallSettings settings) {
     throw new UnsupportedOperationException("Session resumption is not yet implemented.");
   }
+
+  private static <RequestT, ResponseT> UnaryCallable<RequestT, ResponseT> createRetryingCallable(
+      UnaryCallable<RequestT, ResponseT> callable,
+      ResumableUploadCommand command,
+      ClientContext clientContext) {
+    RetryAlgorithm<ResponseT> 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/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

@@ -457,6 +457,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(
@@ -477,11 +573,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;