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 474ab7529bca..b52912ba8ebb 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; @@ -74,6 +76,8 @@ public class ResumableUploadCallableImpl private final ClientContext clientContext; private final UnaryCallable> retryingUploadChunkCallable; + private final UnaryCallable> + retryingQueryCallable; public ResumableUploadCallableImpl( ResumableUploadClient client, @@ -86,6 +90,9 @@ public ResumableUploadCallableImpl( this.retryingUploadChunkCallable = createRetryingCallable( client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD, clientContext); + this.retryingQueryCallable = + createRetryingCallable( + client.queryStatusCallable(), ResumableUploadCommand.QUERY, clientContext); } @Override @@ -109,6 +116,7 @@ public ResumableUploadFuture futureCall( return ResumableUploadFutureImpl.create( startFuture, retryingUploadChunkCallable, + retryingQueryCallable, payload, effectiveSettings, clientContext.getDefaultCallContext()); 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 1b5c62eff3b1..e04deaaa47bc 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,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.ResumableUploadStatus; +import com.google.api.gax.rpc.ResumableUploadErrorClassifier.Category; import com.google.common.util.concurrent.MoreExecutors; import java.io.IOException; import java.io.InputStream; @@ -50,6 +53,9 @@ /** * Coordinates chunk transmission steps of a resumable upload session. * + *

Expects {@code uploadChunkCallable} and {@code queryStatusCallable} to be pre-wrapped in + * retrying callables that handle transient errors. + * * @param the type of the final response message returned once the upload completes */ @InternalApi @@ -62,6 +68,8 @@ final class ResumableUploadChunkCoordinator { private final UnaryCallable> uploadChunkCallable; + private final UnaryCallable> + queryStatusCallable; private final String uploadUrl; private final RewindableStreamBuffer buffer; private final ApiCallContext callContext; @@ -70,12 +78,15 @@ final class ResumableUploadChunkCoordinator { ResumableUploadChunkCoordinator( UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, String uploadUrl, InputStream payload, int chunkSize, ApiCallContext callContext) { 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"); @@ -96,33 +107,23 @@ ApiFuture start() { } private void transmitChunk(long currentOffset) { - // Abort if the session was already completed or canceled. if (result.isDone()) { return; } - - // Read the next chunk slice from the payload stream. try { buffer.fill(currentOffset); - } catch (IOException e) { - result.setException(e); - return; + dispatchCurrentChunk(); + } catch (Throwable t) { + result.setException(t); } + } - // Determine if this is the final chunk and build the chunk request. - ChunkUploadRequest chunkRequest = - ChunkUploadRequest.newBuilder() - .setUploadUrl(uploadUrl) - .setPayload(buffer.getBuffer()) - .setPayloadLength(buffer.getPayloadLength()) - .setOffset(buffer.getBufferBaseOffset()) - .setFinal(buffer.isFinal()) - .build(); - - // Dispatch the chunk upload call and register the in-flight future for cancellation. - long chunkLength = chunkRequest.getPayloadLength(); - boolean isFinal = chunkRequest.isFinal(); + private void dispatchCurrentChunk() { + if (result.isDone()) { + return; + } try { + ChunkUploadRequest chunkRequest = buildCurrentChunkRequest(); ApiFuture> chunkFuture = uploadChunkCallable.futureCall(chunkRequest, callContext); this.currentChunkFuture = chunkFuture; @@ -130,26 +131,61 @@ private void transmitChunk(long currentOffset) { chunkFuture.cancel(true); return; } - ApiFutures.addCallback( chunkFuture, new ApiFutureCallback>() { @Override public void onSuccess(ChunkUploadResponse response) { - if (result.isDone()) { + if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) { + recover(); + } else { + onChunkUploaded(response); + } + } + + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException || result.isDone()) { return; } - long nextOffset = currentOffset + chunkLength; - if (response.getUploadStatus() == ResumableUploadStatus.FINAL) { - result.set(response.getResponse()); - } else if (isFinal) { - result.setException( - new IllegalStateException( - "Upload stream ended and final chunk was transmitted, but server returned" - + " incomplete status for upload URL: " - + uploadUrl)); + Category category = + ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD); + if (category == Category.RECOVERABLE) { + recover(); } else { - chunkExecutor.execute(() -> transmitChunk(nextOffset)); + // Category.TRANSIENT errors reaching here have already exhausted their retry budget + // in the underlying RetryingCallable and become fatal per protocol specification. + result.setException(t); + } + } + }, + chunkExecutor); + } catch (Throwable t) { + result.setException(t); + } + } + + private void recover() { + if (result.isDone()) { + return; + } + try { + ApiFuture> queryFuture = + queryStatusCallable.futureCall(QueryStatusRequest.create(uploadUrl), callContext); + this.currentChunkFuture = queryFuture; + if (result.isCancelled()) { + queryFuture.cancel(true); + return; + } + ApiFutures.addCallback( + queryFuture, + new ApiFutureCallback>() { + @Override + public void onSuccess(QueryStatusResponse queryResponse) { + try { + handleQueryResponse(queryResponse); + } catch (Throwable t) { + result.setException(t); } } @@ -161,10 +197,78 @@ public void onFailure(Throwable t) { result.setException(t); } }, - MoreExecutors.directExecutor()); + chunkExecutor); } catch (Throwable t) { result.setException(t); } } -} + private void handleQueryResponse(QueryStatusResponse queryResponse) throws IOException { + if (result.isDone()) { + return; + } + if (queryResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) { + throw protocolViolation( + "Query status response missing X-Goog-Upload-Status header for upload URL: " + uploadUrl); + } + if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) { + onChunkUploaded( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, queryResponse.getResponse())); + return; + } + Long committedOffset = queryResponse.getCommittedOffset(); + if (committedOffset == null) { + throw protocolViolation( + "Incomplete query status response did not include a committed offset for upload URL: " + + uploadUrl); + } + buffer.realignTo(committedOffset); + dispatchCurrentChunk(); + } + + private void onChunkUploaded(ChunkUploadResponse response) { + if (result.isDone()) { + return; + } + long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength(); + if (response.getUploadStatus() == ResumableUploadStatus.FINAL) { + result.set(response.getResponse()); + } else if (buffer.isFinal()) { + result.setException( + new IllegalStateException( + "Upload stream ended and final chunk was transmitted, but server returned" + + " incomplete status for upload URL: " + + uploadUrl)); + } else { + chunkExecutor.execute(() -> transmitChunk(nextOffset)); + } + } + + private ChunkUploadRequest buildCurrentChunkRequest() { + // Determine if this is the final chunk and build the chunk request. + return ChunkUploadRequest.newBuilder() + .setUploadUrl(uploadUrl) + .setPayload(buffer.getBuffer()) + .setPayloadLength(buffer.getPayloadLength()) + .setOffset(buffer.getBufferBaseOffset()) + .setFinal(buffer.isFinal()) + .build(); + } + + private static final StatusCode FAILED_PRECONDITION_STATUS_CODE = + new StatusCode() { + @Override + public StatusCode.Code getCode() { + return StatusCode.Code.FAILED_PRECONDITION; + } + + @Override + public @Nullable Object getTransportCode() { + return null; + } + }; + + private static FailedPreconditionException protocolViolation(String message) { + return new FailedPreconditionException(message, null, FAILED_PRECONDITION_STATUS_CODE, false); + } +} 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 43b86c084a02..84c257cc4fd0 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,6 +38,8 @@ 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.common.util.concurrent.MoreExecutors; import com.google.errorprone.annotations.concurrent.GuardedBy; @@ -65,6 +67,8 @@ 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; @@ -85,12 +89,18 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur static ResumableUploadFutureImpl create( ApiFuture startFuture, UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, ApiCallContext callContext) { ResumableUploadFutureImpl future = new ResumableUploadFutureImpl<>( - startFuture, uploadChunkCallable, payload, settings, callContext); + startFuture, + uploadChunkCallable, + queryStatusCallable, + payload, + settings, + callContext); try { future.start(); } catch (Throwable t) { @@ -102,12 +112,15 @@ static ResumableUploadFutureImpl create( private ResumableUploadFutureImpl( ApiFuture startFuture, UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, ApiCallContext callContext) { 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"); @@ -128,6 +141,7 @@ public void onSuccess(ResumableUploadSession session) { ResumableUploadChunkCoordinator coordinator = new ResumableUploadChunkCoordinator<>( uploadChunkCallable, + queryStatusCallable, uploadSessionUrl, payload, settings.getChunkSize(), 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 eac9113de782..aa6eb56e0237 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 @@ -44,6 +44,8 @@ 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; @@ -58,6 +60,7 @@ import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; +import org.jspecify.annotations.Nullable; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -70,6 +73,7 @@ class ResumableUploadCallableImplTest { private ResumableUploadClient mockClient; private UnaryCallable mockStartCallable; private UnaryCallable> mockChunkCallable; + private UnaryCallable> mockQueryCallable; private ResumableUploadCallSettings defaultSettings; private FakeCallContext callContext; @@ -81,9 +85,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(); @@ -478,6 +484,264 @@ 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(FailedPreconditionException.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(FailedPreconditionException.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(FailedPreconditionException.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()); + } + private static class HttpStatusStatusCode implements StatusCode { private final int httpStatus; private final StatusCode.Code code; @@ -503,6 +767,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(