From 08e807a96f060ab3af093a20f5b824c7558a65c7 Mon Sep 17 00:00:00 2001 From: whowes Date: Sat, 12 Sep 2026 06:14:03 +0000 Subject: [PATCH] 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. --- .../gax/rpc/ResumableUploadCallableImpl.java | 36 +++++- .../rpc/ResumableUploadChunkCoordinator.java | 12 +- .../com/google/api/gax/rpc/CallableTest.java | 2 + .../rpc/ResumableUploadCallableImplTest.java | 106 ++++++++++++++++++ 4 files changed, 152 insertions(+), 4 deletions(-) diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java index 318cdf308d51..474ab7529bca 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java +++ b/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 extends ResumableUploadCallable { + 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 client; private final ResumableUploadCallSettings defaultCallSettings; private final ClientContext clientContext; + private final UnaryCallable> + retryingUploadChunkCallable; public ResumableUploadCallableImpl( ResumableUploadClient client, @@ -66,6 +83,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 +108,7 @@ public ResumableUploadFuture futureCall( return ResumableUploadFutureImpl.create( startFuture, - client.uploadChunkCallable(), + retryingUploadChunkCallable, payload, effectiveSettings, clientContext.getDefaultCallContext()); @@ -99,4 +119,18 @@ public ResumableUploadFuture resumeCall( String sessionUrl, InputStream payload, @Nullable ResumableUploadCallSettings settings) { throw new UnsupportedOperationException("Session resumption is not yet implemented."); } + + private static UnaryCallable createRetryingCallable( + UnaryCallable callable, + ResumableUploadCommand command, + ClientContext clientContext) { + RetryAlgorithm 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())); + } } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java index fea348f9effc..69e7318aa8fa 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java +++ b/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,10 @@ final class ResumableUploadChunkCoordinator { private static final byte[] EMPTY_PAYLOAD = new byte[0]; + // Serializes buffer mutations across multiple threads (i.e. from retry/recovery) + private final Executor chunkExecutor = + MoreExecutors.newSequentialExecutor(MoreExecutors.directExecutor()); + private final UnaryCallable> uploadChunkCallable; private final String uploadUrl; @@ -93,7 +98,7 @@ ApiFuture start() { } }, MoreExecutors.directExecutor()); - transmitChunk(0L); + chunkExecutor.execute(() -> transmitChunk(0L)); return result; } @@ -157,9 +162,10 @@ public void onSuccess(ChunkUploadResponse 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)); } } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java index 745cb294bb7b..4a905d50bbd3 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java +++ b/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 uploadClient = mock(ResumableUploadClient.class, Mockito.withSettings().withoutAnnotations()); + when(uploadClient.uploadChunkCallable()) + .thenReturn(mock(UnaryCallable.class, Mockito.withSettings().withoutAnnotations())); ResumableUploadCallSettings settings = ResumableUploadCallSettings.newBuilder().setChunkSize(1024).build(); diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java index b882a43df401..01e0d4b5bec0 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java +++ b/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 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 captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any()); + List 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 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> chunkAttempt0Future = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(chunkAttempt0Future) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "should-not-reach"))); + + ResumableUploadFuture 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;