From 566ca010191b7ce611150b7febb8d884da974a91 Mon Sep 17 00:00:00 2001 From: whowes Date: Sat, 12 Sep 2026 15:58:07 +0000 Subject: [PATCH] feat(gax): add recoverable error query and buffer realignment loop Introduces a query-and-realign loop when encountering recoverable protocol errors during chunk uploads. Queries the server for the committed offset and adjusts the buffer window before resuming chunk transmission. --- .../gax/rpc/ResumableUploadCallableImpl.java | 13 +- .../rpc/ResumableUploadChunkCoordinator.java | 55 ++-- .../gax/rpc/ResumableUploadFutureImpl.java | 16 +- .../com/google/api/gax/rpc/CallableTest.java | 2 + .../rpc/ResumableUploadCallableImplTest.java | 275 ++++++++++++++++++ 5 files changed, 330 insertions(+), 31 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 7d12555585ed..ca6ff1bbbe2a 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, @@ -89,6 +93,8 @@ public ResumableUploadCallableImpl( .build(); this.retryingUploadChunkCallable = createRetryingCallable(client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD); + this.retryingQueryCallable = + createRetryingCallable(client.queryStatusCallable(), ResumableUploadCommand.QUERY); } @Override @@ -110,7 +116,12 @@ public ResumableUploadFuture futureCall( } return ResumableUploadFutureImpl.create( - startFuture, retryingUploadChunkCallable, payload, effectiveSettings, clientContext); + startFuture, + retryingUploadChunkCallable, + retryingQueryCallable, + payload, + effectiveSettings, + clientContext); } @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 c24078bd513c..ac33ca668f2b 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 @@ -210,24 +210,22 @@ private void onChunkUploaded(ChunkUploadResponse response) { if (result.isDone()) { return; } - try { - if (response.getUploadStatus() == ResumableUploadStatus.FINAL) { - result.set(response.getResponse()); - } else if (buffer.isFinal()) { - throw new IllegalStateException( - "Upload stream ended and final chunk was transmitted, but server returned incomplete" - + " status for upload URL: " - + uploadUrl); - } else { - long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength(); - transmitChunk(nextOffset); - } - } catch (Throwable t) { - result.setException(t); + 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()) @@ -237,21 +235,20 @@ private ChunkUploadRequest buildCurrentChunkRequest() { .build(); } - private static FailedPreconditionException protocolViolation(String message) { - return new FailedPreconditionException( - message, - null, - new StatusCode() { - @Override - public StatusCode.Code getCode() { - return StatusCode.Code.FAILED_PRECONDITION; - } + 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; - } - }, - false); + @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 3dd032085e10..4662813d0844 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 ClientContext clientContext; @@ -85,12 +89,18 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur static ResumableUploadFutureImpl create( ApiFuture startFuture, UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, ClientContext clientContext) { ResumableUploadFutureImpl future = new ResumableUploadFutureImpl<>( - startFuture, uploadChunkCallable, payload, settings, clientContext); + startFuture, + uploadChunkCallable, + queryStatusCallable, + payload, + settings, + clientContext); 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, ClientContext clientContext) { 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(