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 4d66f5539c9f..25c5e96fb0d6 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 @@ -37,6 +37,8 @@ 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.QueryStatusRequest; +import com.google.api.gax.resumable.QueryStatusResponse; import com.google.api.gax.resumable.ResumableUploadClient; import com.google.api.gax.resumable.ResumableUploadSession; import com.google.api.gax.retrying.ExponentialRetryAlgorithm; @@ -78,6 +80,9 @@ public class ResumableUploadCallableImpl private final ClientContext clientContext; private final UnaryCallable> retryingUploadChunkCallable; + private final UnaryCallable> + retryingQueryCallable; + private final ExponentialRetryAlgorithm recoveryAlgorithm; public ResumableUploadCallableImpl( ResumableUploadClient client, @@ -90,6 +95,11 @@ public ResumableUploadCallableImpl( this.retryingUploadChunkCallable = createRetryingCallable( client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD, clientContext); + this.retryingQueryCallable = + createRetryingCallable( + client.queryStatusCallable(), ResumableUploadCommand.QUERY, clientContext); + this.recoveryAlgorithm = + new ExponentialRetryAlgorithm(RETRY_SETTINGS, clientContext.getClock()); } @Override @@ -113,9 +123,11 @@ public ResumableUploadFuture futureCall( return ResumableUploadFutureImpl.create( startFuture, retryingUploadChunkCallable, + retryingQueryCallable, payload, effectiveSettings, - clientContext.getDefaultCallContext()); + clientContext, + recoveryAlgorithm); } @Override 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 116eb589050b..855b6037f018 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 @@ -38,10 +38,19 @@ import com.google.api.core.SettableApiFuture; import com.google.api.gax.resumable.ChunkUploadRequest; import com.google.api.gax.resumable.ChunkUploadResponse; +import com.google.api.gax.resumable.QueryStatusRequest; +import com.google.api.gax.resumable.QueryStatusResponse; import com.google.api.gax.resumable.ResumableUploadStatus; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; +import com.google.api.gax.retrying.TimedAttemptSettings; +import com.google.api.gax.rpc.ResumableUploadErrorClassifier.Category; import com.google.common.util.concurrent.MoreExecutors; +import java.io.IOException; import java.io.InputStream; import java.util.concurrent.CancellationException; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -56,23 +65,37 @@ final class ResumableUploadChunkCoordinator { private final UnaryCallable> uploadChunkCallable; + private final UnaryCallable> + queryStatusCallable; private final String uploadUrl; private final RewindableStreamBuffer buffer; private final ApiCallContext callContext; + private final ExponentialRetryAlgorithm recoveryAlgorithm; + private final ScheduledExecutorService executor; private final SettableApiFuture uploadResultFuture = SettableApiFuture.create(); - private volatile @Nullable ApiFuture inFlightFuture; + private volatile @Nullable Future inFlightFuture; + private TimedAttemptSettings recoverySettings; + private boolean madeProgressSinceRecovery = true; ResumableUploadChunkCoordinator( UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, String uploadUrl, InputStream payload, int chunkSize, - ApiCallContext callContext) { + ApiCallContext callContext, + ExponentialRetryAlgorithm recoveryAlgorithm, + ScheduledExecutorService executor) { this.uploadChunkCallable = checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); + this.queryStatusCallable = + checkNotNull(queryStatusCallable, "queryStatusCallable must not be null"); this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null"); checkNotNull(payload, "payload must not be null"); this.callContext = checkNotNull(callContext, "callContext must not be null"); + this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null"); + this.executor = checkNotNull(executor, "executor must not be null"); + this.recoverySettings = recoveryAlgorithm.createFirstAttempt(); this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl); } @@ -83,13 +106,18 @@ ApiFuture getFuture() { void start() { uploadResultFuture.addListener( () -> { - ApiFuture inFlight = inFlightFuture; + Future inFlight = inFlightFuture; if (uploadResultFuture.isCancelled() && inFlight != null) { inFlight.cancel(true); } }, MoreExecutors.directExecutor()); - transmitChunk(); + try { + buffer.fill(); + transmitChunk(); + } catch (Throwable t) { + uploadResultFuture.setException(t); + } } private void transmitChunk() { @@ -99,26 +127,14 @@ private void transmitChunk() { return; } - // Read the next chunk slice from the payload stream. - buffer.fill(); - - // Determine if this is the final chunk and build the chunk request. - ChunkUploadRequest chunkRequest = - ChunkUploadRequest.newBuilder() - .setUploadUrl(uploadUrl) - .setPayload(buffer.getPayload()) - .setOffset(buffer.getBufferBaseOffset()) - .setFinal(buffer.isFinal()) - .build(); + ChunkUploadRequest chunkRequest = buildCurrentChunkRequest(); // Dispatch the chunk upload call and register the in-flight future for cancellation. - boolean isFinal = chunkRequest.isFinal(); ApiFuture> chunkFuture = uploadChunkCallable.futureCall(chunkRequest, callContext); if (!tryRegisterInFlightFuture(chunkFuture)) { return; } - ApiFutures.addCallback( chunkFuture, new ApiFutureCallback>() { @@ -127,15 +143,88 @@ public void onSuccess(ChunkUploadResponse response) { if (uploadResultFuture.isDone()) { return; } - if (response.getUploadStatus() == ResumableUploadStatus.FINAL) { - uploadResultFuture.set(response.getResponse()); - } else if (isFinal) { - uploadResultFuture.setException( + if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) { + recover( new IllegalStateException( - "Upload stream ended and final chunk was transmitted, but server returned" - + " incomplete status")); + "Chunk upload response missing X-Goog-Upload-Status header for upload URL: " + + uploadUrl)); } else { - transmitChunk(); + try { + handleChunkResponse(response); + } catch (Throwable t) { + uploadResultFuture.setException(t); + } + } + } + + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException || uploadResultFuture.isDone()) { + return; + } + Category category = + ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD); + if (category == Category.RECOVERABLE) { + recover(t); + } else { + // Category.TRANSIENT errors reaching here have already exhausted their retry budget + // in the underlying RetryingCallable and become fatal per protocol specification. + uploadResultFuture.setException(t); + } + } + }, + MoreExecutors.directExecutor()); + } catch (Throwable t) { + uploadResultFuture.setException(t); + } + } + + private void recover(Throwable cause) { + try { + if (madeProgressSinceRecovery) { + recoverySettings = + recoveryAlgorithm.createNextAttempt(recoveryAlgorithm.createFirstAttempt()); + } else { + recoverySettings = recoveryAlgorithm.createNextAttempt(recoverySettings); + } + madeProgressSinceRecovery = false; + if (!recoveryAlgorithm.shouldRetry(recoverySettings)) { + uploadResultFuture.setException(cause); + return; + } + tryRegisterInFlightFuture( + executor.schedule( + this::queryStatus, + recoverySettings.getRandomizedRetryDelayDuration().toNanos(), + TimeUnit.NANOSECONDS)); + } catch (Throwable t) { + uploadResultFuture.setException(t); + } + } + + private void queryStatus() { + try { + if (uploadResultFuture.isDone()) { + return; + } + // Dispatch the query status call and register the in-flight future for cancellation. + ApiFuture> queryFuture = + queryStatusCallable.futureCall(QueryStatusRequest.create(uploadUrl), callContext); + if (!tryRegisterInFlightFuture(queryFuture)) { + return; + } + ApiFutures.addCallback( + queryFuture, + new ApiFutureCallback>() { + @Override + public void onSuccess(QueryStatusResponse queryResponse) { + if (uploadResultFuture.isDone()) { + return; + } + try { + handleQueryResponse(queryResponse); + } catch (Throwable t) { + uploadResultFuture.setException(t); } } @@ -157,7 +246,7 @@ public void onFailure(Throwable t) { * Registers the in-flight future for possible cancellation, returning false if the upload was * already cancelled. */ - private boolean tryRegisterInFlightFuture(ApiFuture future) { + private boolean tryRegisterInFlightFuture(Future future) { this.inFlightFuture = future; if (uploadResultFuture.isCancelled()) { future.cancel(true); @@ -165,4 +254,50 @@ private boolean tryRegisterInFlightFuture(ApiFuture future) { } return true; } + + private void handleQueryResponse(QueryStatusResponse queryResponse) + throws IOException { + if (queryResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) { + throw new IllegalStateException( + "Query status response missing X-Goog-Upload-Status header for upload URL: " + uploadUrl); + } + if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) { + uploadResultFuture.set(queryResponse.getResponse()); + return; + } + Long committedOffset = queryResponse.getCommittedOffset(); + if (committedOffset == null) { + throw new IllegalStateException( + "Incomplete query status response did not include a committed offset for upload URL: " + + uploadUrl); + } + buffer.realignTo(committedOffset); + transmitChunk(); + } + + private void handleChunkResponse(ChunkUploadResponse response) throws IOException { + if (response.getUploadStatus() == ResumableUploadStatus.FINAL) { + uploadResultFuture.set(response.getResponse()); + } else if (buffer.isFinal()) { + uploadResultFuture.setException( + new IllegalStateException( + "Upload stream ended and final chunk was transmitted, but server returned" + + " incomplete status for upload URL: " + + uploadUrl)); + } else { + madeProgressSinceRecovery = true; + buffer.fill(); + transmitChunk(); + } + } + + private ChunkUploadRequest buildCurrentChunkRequest() { + // Determine if this is the final chunk and build the chunk request. + return ChunkUploadRequest.newBuilder() + .setUploadUrl(uploadUrl) + .setPayload(buffer.getPayload()) + .setOffset(buffer.getBufferBaseOffset()) + .setFinal(buffer.isFinal()) + .build(); + } } diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java index 6b4c640f70bf..451760fe2c14 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java @@ -38,7 +38,10 @@ import com.google.api.core.SettableApiFuture; import com.google.api.gax.resumable.ChunkUploadRequest; import com.google.api.gax.resumable.ChunkUploadResponse; +import com.google.api.gax.resumable.QueryStatusRequest; +import com.google.api.gax.resumable.QueryStatusResponse; import com.google.api.gax.resumable.ResumableUploadSession; +import com.google.api.gax.retrying.ExponentialRetryAlgorithm; import com.google.common.util.concurrent.MoreExecutors; import com.google.errorprone.annotations.concurrent.GuardedBy; import java.io.IOException; @@ -46,6 +49,7 @@ import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import org.jspecify.annotations.NullMarked; @@ -65,9 +69,13 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur private final ApiFuture startFuture; private final UnaryCallable> uploadChunkCallable; + private final UnaryCallable> + queryStatusCallable; private final InputStream payload; private final ResumableUploadCallSettings settings; private final ApiCallContext callContext; + private final ScheduledExecutorService executor; + private final ExponentialRetryAlgorithm recoveryAlgorithm; private final SettableApiFuture resultFuture = SettableApiFuture.create(); private volatile @Nullable String uploadSessionUrl; @@ -86,12 +94,20 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur static ResumableUploadFutureImpl create( ApiFuture startFuture, UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, - ApiCallContext callContext) { + ClientContext clientContext, + ExponentialRetryAlgorithm recoveryAlgorithm) { ResumableUploadFutureImpl future = new ResumableUploadFutureImpl<>( - startFuture, uploadChunkCallable, payload, settings, callContext); + startFuture, + uploadChunkCallable, + queryStatusCallable, + payload, + settings, + clientContext, + recoveryAlgorithm); try { future.start(); } catch (Throwable t) { @@ -103,16 +119,23 @@ static ResumableUploadFutureImpl create( private ResumableUploadFutureImpl( ApiFuture startFuture, UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, - ApiCallContext callContext) { + ClientContext clientContext, + ExponentialRetryAlgorithm recoveryAlgorithm) { this.startFuture = checkNotNull(startFuture, "startFuture must not be null"); this.uploadChunkCallable = checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); + this.queryStatusCallable = + checkNotNull(queryStatusCallable, "queryStatusCallable 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"); + checkNotNull(clientContext, "clientContext must not be null"); + this.callContext = clientContext.getDefaultCallContext(); + this.executor = checkNotNull(clientContext.getExecutor(), "executor must not be null"); + this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null"); this.inFlightFuture = startFuture; } @@ -125,7 +148,14 @@ public void onSuccess(ResumableUploadSession session) { String sessionUrl = session.getUploadUrl(); ResumableUploadChunkCoordinator coordinator = new ResumableUploadChunkCoordinator<>( - uploadChunkCallable, sessionUrl, payload, settings.getChunkSize(), callContext); + uploadChunkCallable, + queryStatusCallable, + sessionUrl, + payload, + settings.getChunkSize(), + callContext, + recoveryAlgorithm, + executor); ApiFuture uploadFuture = coordinator.getFuture(); synchronized (lock) { if (resultFuture.isDone()) { 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 4a905d50bbd3..0cf4280a43ca 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 @@ -216,6 +216,8 @@ void testResumableUploadCallable() { mock(ResumableUploadClient.class, Mockito.withSettings().withoutAnnotations()); when(uploadClient.uploadChunkCallable()) .thenReturn(mock(UnaryCallable.class, Mockito.withSettings().withoutAnnotations())); + when(uploadClient.queryStatusCallable()) + .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 4f3c077ef1ce..ad6062318b2c 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 @@ -40,11 +40,14 @@ import static org.mockito.Mockito.when; import static org.mockito.Mockito.withSettings; +import com.google.api.core.ApiFuture; import com.google.api.core.ApiFutures; import com.google.api.core.ForwardingApiFuture; import com.google.api.core.SettableApiFuture; import com.google.api.gax.resumable.ChunkUploadRequest; import com.google.api.gax.resumable.ChunkUploadResponse; +import com.google.api.gax.resumable.QueryStatusRequest; +import com.google.api.gax.resumable.QueryStatusResponse; import com.google.api.gax.resumable.ResumableUploadClient; import com.google.api.gax.resumable.ResumableUploadSession; import com.google.api.gax.resumable.ResumableUploadStatus; @@ -59,6 +62,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; +import org.jspecify.annotations.Nullable; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -71,6 +75,7 @@ class ResumableUploadCallableImplTest { private ResumableUploadClient mockClient; private UnaryCallable mockStartCallable; private UnaryCallable> mockChunkCallable; + private UnaryCallable> mockQueryCallable; private ResumableUploadCallSettings defaultSettings; private FakeCallContext callContext; @@ -82,9 +87,11 @@ void setUp() { mockClient = mock(ResumableUploadClient.class, withSettings().withoutAnnotations()); mockStartCallable = mock(UnaryCallable.class, withSettings().withoutAnnotations()); mockChunkCallable = mock(UnaryCallable.class, withSettings().withoutAnnotations()); + mockQueryCallable = mock(UnaryCallable.class, withSettings().withoutAnnotations()); lenient().when(mockClient.startUploadCallable()).thenReturn(mockStartCallable); lenient().when(mockClient.uploadChunkCallable()).thenReturn(mockChunkCallable); + lenient().when(mockClient.queryStatusCallable()).thenReturn(mockQueryCallable); defaultSettings = ResumableUploadCallSettings.newBuilder().setChunkSize(8).build(); callContext = FakeCallContext.createDefault(); @@ -528,6 +535,331 @@ void testChunkRetry_cancellationDuringBackoff_deschedulesPendingAttempt() { verify(mockChunkCallable, times(1)).futureCall(any(), any()); } + @Test + void testRecovery_recoverableChunkError_recoversViaQuery() throws Exception { + stubStartSession("https://upload.url/recovery-success"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.newBuilder() + .setUploadStatus(ResumableUploadStatus.FINAL) + .setResponse("recovered-response") + .build())); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(0L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("recovered-response"); + assertThat(future.isDone()).isTrue(); + verify(mockQueryCallable, times(1)).futureCall(any(), any()); + verify(mockChunkCallable, times(2)).futureCall(any(), any()); + } + + @Test + void testRecovery_queryReturnsFinal_completesWithoutResending() throws Exception { + stubStartSession("https://upload.url/recovery-already-complete"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(5L, "server-finalized", ResumableUploadStatus.FINAL))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("server-finalized"); + assertThat(future.isDone()).isTrue(); + verify(mockQueryCallable, times(1)).futureCall(any(), any()); + verify(mockChunkCallable, times(1)).futureCall(any(), any()); + } + + @Test + void testRecovery_queryNullOffset_failsFatal() throws Exception { + stubStartSession("https://upload.url/recovery-null-offset"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(null, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(exception.getCause()) + .hasMessageThat() + .contains("did not include a committed offset"); + verify(mockQueryCallable, times(1)).futureCall(any(), any()); + verify(mockChunkCallable, times(1)).futureCall(any(), any()); + } + + @Test + void testRecovery_committedMidBuffer_compactsAndTopsUp() throws Exception { + stubStartSession("https://upload.url/recovery-compact-topup"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "all-done"))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(4L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("0123456789abcdef"), null); + + assertThat(future.get()).isEqualTo("all-done"); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(3)).futureCall(captor.capture(), any()); + List requests = captor.getAllValues(); + assertChunk(requests.get(0), 0, 8, false); + assertChunk(requests.get(1), 4, 8, false); + assertChunk(requests.get(2), 12, 4, true); + } + + @Test + void testRecovery_offsetBelowBufferBase_failsFatal() throws Exception { + stubStartSession("https://upload.url/recovery-below-base"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null))) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(4L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("0123456789abcdef"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(exception.getCause()).hasMessageThat().contains("below buffer base offset"); + } + + @Test + void testRecovery_missingStatusHeaderOn200_triggersRecovery() throws Exception { + stubStartSession("https://upload.url/missing-status-200"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.UNKNOWN, null))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "recovered-ok"))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(0L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("recovered-ok"); + verify(mockQueryCallable, times(1)).futureCall(any(), any()); + verify(mockChunkCallable, times(2)).futureCall(any(), any()); + } + + @Test + void testRecovery_finalChunk_preservesFinalFlag() throws Exception { + stubStartSession("https://upload.url/recovery-final-chunk"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null))) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "final-chunk-done"))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(10L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("0123456789ab"), null); + + assertThat(future.get()).isEqualTo("final-chunk-done"); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(3)).futureCall(captor.capture(), any()); + List requests = captor.getAllValues(); + assertChunk(requests.get(0), 0, 8, false); + assertChunk(requests.get(1), 8, 4, true); + assertChunk(requests.get(2), 10, 2, true); + } + + @Test + void testRecovery_queryMissingStatusHeader_failsFatal() throws Exception { + stubStartSession("https://upload.url/recovery-query-missing-status"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(0L, null, ResumableUploadStatus.UNKNOWN))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(exception.getCause()) + .hasMessageThat() + .contains("missing X-Goog-Upload-Status header"); + } + + @Test + void testRecovery_queryRecoverableError_failsFatal() throws Exception { + stubStartSession("https://upload.url/recovery-query-cat2"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + // 400 is recoverable for UPLOAD, but fatal for QUERY + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(ApiException.class); + verify(mockQueryCallable, times(1)).futureCall(any(), any()); + } + + @Test + void testRecovery_queryTransientError_retriesAndSucceeds() throws Exception { + stubStartSession("https://upload.url/recovery-query-transient"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "query-retry-ok"))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture(createApiException(503, StatusCode.Code.UNAVAILABLE))) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(0L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("query-retry-ok"); + verify(mockQueryCallable, times(2)).futureCall(any(), any()); + verify(mockChunkCallable, times(2)).futureCall(any(), any()); + } + + @Test + void testRecovery_repeatedWithoutProgress_failsAfterMaxAttempts() throws Exception { + stubStartSession("https://upload.url/recovery-exhausted"); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(0L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(ApiException.class); + verify(mockQueryCallable, times(4)).futureCall(any(), any()); + verify(mockChunkCallable, times(5)).futureCall(any(), any()); + } + + @Test + void testRecovery_progressBetweenRecoveries_resetsAttempts() throws Exception { + stubStartSession("https://upload.url/recovery-reset"); + ApiFuture> recoverable = + ApiFutures.immediateFailedFuture(createApiException(400, StatusCode.Code.INVALID_ARGUMENT)); + ApiFuture> active = + ApiFutures.immediateFuture(ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null)); + when(mockChunkCallable.futureCall(any(), any())) + .thenReturn(recoverable) + .thenReturn(active) + .thenReturn(recoverable) + .thenReturn(active) + .thenReturn(recoverable) + .thenReturn(active) + .thenReturn(recoverable) + .thenReturn(active) + .thenReturn(recoverable) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "reset-ok"))); + + when(mockQueryCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture(createQueryResponse(0L, null, ResumableUploadStatus.ACTIVE))) + .thenReturn( + ApiFutures.immediateFuture(createQueryResponse(8L, null, ResumableUploadStatus.ACTIVE))) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(16L, null, ResumableUploadStatus.ACTIVE))) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(24L, null, ResumableUploadStatus.ACTIVE))) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(32L, null, ResumableUploadStatus.ACTIVE))); + + ResumableUploadFuture future = + callable.futureCall( + "resource-path", streamOf("0123456789abcdefghijklmnopqrstuvwxyz0123"), null); + + assertThat(future.get()).isEqualTo("reset-ok"); + verify(mockQueryCallable, times(5)).futureCall(any(), any()); + verify(mockChunkCallable, times(10)).futureCall(any(), any()); + } + private static class HttpStatusStatusCode implements StatusCode { private final int httpStatus; private final StatusCode.Code code; @@ -553,6 +885,17 @@ private static ApiException createApiException(int httpStatus, StatusCode.Code c "HTTP " + httpStatus, null, new HttpStatusStatusCode(httpStatus, code), false); } + private static QueryStatusResponse createQueryResponse( + @Nullable Long committedOffset, + @Nullable String response, + ResumableUploadStatus uploadStatus) { + return QueryStatusResponse.newBuilder() + .setCommittedOffset(committedOffset) + .setResponse(response) + .setUploadStatus(uploadStatus) + .build(); + } + private void stubStartSession(String uploadUrl) { when(mockStartCallable.futureCall(any(), any())) .thenReturn(