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..dbb63b0184cf 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 @@ -91,7 +91,8 @@ public ResumableUploadFuture futureCall( client.uploadChunkCallable(), payload, effectiveSettings, - clientContext.getDefaultCallContext()); + clientContext.getDefaultCallContext(), + 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 0703b16a9b7b..64b167d99538 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 @@ -39,13 +39,21 @@ import com.google.api.gax.resumable.ChunkUploadRequest; import com.google.api.gax.resumable.ChunkUploadResponse; 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.RetryingFuture; +import com.google.api.gax.retrying.ScheduledRetryingExecutor; import com.google.common.io.ByteStreams; import com.google.common.util.concurrent.MoreExecutors; import com.google.errorprone.annotations.concurrent.GuardedBy; import java.io.IOException; import java.io.InputStream; +import java.time.Duration; import java.util.Arrays; import java.util.concurrent.CancellationException; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicLong; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -57,18 +65,34 @@ @NullMarked final class ResumableUploadChunkCoordinator { + static final RetrySettings DEFAULT_CHUNK_RETRY_SETTINGS = + RetrySettings.newBuilder() + .setInitialRetryDelayDuration(Duration.ofMillis(100)) + .setRetryDelayMultiplier(1.3) + .setMaxRetryDelayDuration(Duration.ofMinutes(1)) + .setInitialRpcTimeoutDuration(Duration.ofSeconds(30)) + .setRpcTimeoutMultiplier(1.0) + .setMaxRpcTimeoutDuration(Duration.ofSeconds(30)) + .setTotalTimeoutDuration(Duration.ofMinutes(5)) + .setMaxAttempts(5) + .build(); + private static final byte[] EMPTY_PAYLOAD = new byte[0]; private final Object lock = new Object(); + private final AtomicBoolean dispatching = new AtomicBoolean(false); + private final AtomicLong nextChunkOffset = new AtomicLong(-1L); private final SettableApiFuture result; private final ApiFuture startFuture; - private final UnaryCallable> - uploadChunkCallable; + private final RetryingCallable> + retryingChunkCallable; private final InputStream payload; private final byte[] buffer; private final int chunkSize; private final ApiCallContext callContext; + private final ClientContext clientContext; + private final RetrySettings chunkRetrySettings; private volatile @Nullable String uploadSessionUrl; @@ -87,17 +111,30 @@ final class ResumableUploadChunkCoordinator { UnaryCallable> uploadChunkCallable, InputStream payload, ResumableUploadCallSettings settings, - ApiCallContext callContext) { + ApiCallContext callContext, + ClientContext clientContext) { this.result = checkNotNull(result, "result must not be null"); this.startFuture = checkNotNull(startFuture, "startFuture must not be null"); - this.uploadChunkCallable = - checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); + checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); this.payload = checkNotNull(payload, "payload must not be null"); checkNotNull(settings, "settings must not be null"); checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0"); this.chunkSize = settings.getChunkSize(); this.callContext = checkNotNull(callContext, "callContext must not be null"); + this.clientContext = checkNotNull(clientContext, "clientContext must not be null"); + this.chunkRetrySettings = DEFAULT_CHUNK_RETRY_SETTINGS; this.buffer = new byte[chunkSize]; + + RetryAlgorithm> retryAlgorithm = + new RetryAlgorithm<>( + new ResumableUploadResultRetryAlgorithm<>(ResumableUploadCommand.UPLOAD), + new ExponentialRetryAlgorithm(chunkRetrySettings, clientContext.getClock())); + ScheduledRetryingExecutor> retryingExecutor = + new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor()); + this.retryingChunkCallable = + new RetryingCallable<>( + clientContext.getDefaultCallContext(), uploadChunkCallable, retryingExecutor); + synchronized (lock) { this.inFlightFuture = startFuture; } @@ -115,7 +152,7 @@ public void onSuccess(ResumableUploadSession session) { } } uploadSessionUrl = session.getUploadUrl(); - transmitChunk(0L); + scheduleNextChunk(0L); } @Override @@ -197,7 +234,24 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) { } } - private void transmitChunk(long currentOffset) { + private void scheduleNextChunk(long offset) { + nextChunkOffset.set(offset); + if (dispatching.compareAndSet(false, true)) { + driveLoop(); + } + } + + private void driveLoop() { + do { + long offset = nextChunkOffset.getAndSet(-1L); + if (offset >= 0) { + transmitSingleChunk(offset); + } + dispatching.set(false); + } while (nextChunkOffset.get() >= 0 && dispatching.compareAndSet(false, true)); + } + + private void transmitSingleChunk(long currentOffset) { synchronized (lock) { if (done) { return; @@ -236,47 +290,44 @@ private void transmitChunk(long currentOffset) { .setFinal(isFinal) .build(); - long chunkLength = chunkPayload.length; - try { - ApiFuture> chunkFuture = - uploadChunkCallable.futureCall(chunkRequest, callContext); - setInFlightFuture(chunkFuture); + RetryingFuture> retryingFuture = + retryingChunkCallable.futureCall(chunkRequest, callContext); + setInFlightFuture(retryingFuture); - ApiFutures.addCallback( - chunkFuture, - new ApiFutureCallback>() { - @Override - public void onSuccess(ChunkUploadResponse response) { - synchronized (lock) { - if (done) { - return; - } - } - long nextOffset = currentOffset + chunkLength; - if (response.isComplete()) { - finish(response.getResponse(), null); - } else if (isFinal) { - finish( - null, - new IllegalStateException( - "Upload stream ended and final chunk was transmitted, but server returned" - + " incomplete status")); - } else { - transmitChunk(nextOffset); + long chunkLength = chunkPayload.length; + ApiFutures.addCallback( + retryingFuture, + new ApiFutureCallback>() { + @Override + public void onSuccess(ChunkUploadResponse response) { + synchronized (lock) { + if (done) { + return; } } + long nextOffset = currentOffset + chunkLength; + if (response.isComplete()) { + finish(response.getResponse(), null); + } else if (isFinal) { + finish( + null, + new IllegalStateException( + "Upload stream ended and final chunk was transmitted, but server returned" + + " incomplete status for upload URL: " + + url)); + } else { + scheduleNextChunk(nextOffset); + } + } - @Override - public void onFailure(Throwable t) { - if (t instanceof CancellationException) { - return; - } - finish(null, t); + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException) { + return; } - }, - MoreExecutors.directExecutor()); - } catch (Throwable t) { - finish(null, t); - } + finish(null, t); + } + }, + MoreExecutors.directExecutor()); } } 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 094727707e07..02435da45690 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 @@ -61,10 +61,32 @@ static ResumableUploadFutureImpl create( InputStream payload, ResumableUploadCallSettings settings, ApiCallContext callContext) { + return create( + startFuture, + uploadChunkCallable, + payload, + settings, + callContext, + ClientContext.newBuilder().setDefaultCallContext(callContext).build()); + } + + static ResumableUploadFutureImpl create( + ApiFuture startFuture, + UnaryCallable> uploadChunkCallable, + InputStream payload, + ResumableUploadCallSettings settings, + ApiCallContext callContext, + ClientContext clientContext) { SettableApiFuture result = SettableApiFuture.create(); ResumableUploadChunkCoordinator coordinator = new ResumableUploadChunkCoordinator<>( - result, startFuture, uploadChunkCallable, payload, settings, callContext); + result, + startFuture, + uploadChunkCallable, + payload, + settings, + callContext, + clientContext); ResumableUploadFutureImpl handle = new ResumableUploadFutureImpl<>(result, coordinator); coordinator.start(); 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 3654796f83fe..469ab676fbb3 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 @@ -347,6 +347,225 @@ void testResumeCall_throwsUnsupportedOperationException() { () -> callable.resumeCall("https://upload.url/session", streamOf("data"), null)); } + @Test + void testChunkRetry_cat1FailureThenSuccess_retriesAndSucceeds() throws Exception { + stubStartSession("https://upload.url/chunk-retry-ok"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFailedFuture(createApiException(503, StatusCode.Code.UNAVAILABLE))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "chunk-done"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("chunk-done"); + assertThat(future.isDone()).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("hello".getBytes(StandardCharsets.UTF_8)); + assertThat(requests.get(1).getOffset()).isEqualTo(0); + assertThat(requests.get(1).getPayload()).isEqualTo("hello".getBytes(StandardCharsets.UTF_8)); + } + + @Test + void testChunkRetry_streamNotAdvancedByRetry_sameBytesSent() throws Exception { + stubStartSession("https://upload.url/stream-no-advance"); + ByteCountingStream stream = new ByteCountingStream("01234567"); // exactly 1 chunk of 8 bytes + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + // Chunk 0, attempt 0 -> fails with 503 + .thenReturn( + ApiFutures.immediateFailedFuture(createApiException(503, StatusCode.Code.UNAVAILABLE))) + // Chunk 0, attempt 1 -> succeeds with final response + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "final-response"))); + + ResumableUploadFuture future = callable.futureCall("resource-path", stream, null); + + assertThat(future.get()).isEqualTo("final-response"); + assertThat(future.isDone()).isTrue(); + + // Stream must only have been read exactly 8 bytes (the payload length), + // proving retry did not re-read or advance the input stream. + assertThat(stream.totalBytesRead).isEqualTo(8); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any()); + List requests = captor.getAllValues(); + // Both attempts correspond to chunk 0 and transmit identical bytes + 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_cat2Failure_failsFastWithoutRetrying() { + stubStartSession("https://upload.url/chunk-cat2-fail"); + // HTTP 400 Bad Request is Category 2 (RECOVERABLE) + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), 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); + assertThat(((ApiException) exception.getCause()).getStatusCode().getTransportCode()) + .isEqualTo(400); + + // In G4, Category 2 is not yet retryable / recoverable, so it fails after 1 attempt. + verify(mockChunkCallable, times(1)).futureCall(any(), any()); + } + + @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_cancelSession_cancelsInFlightHttpFuture() throws Exception { + stubStartSession("https://upload.url/cancel-in-flight"); + CountDownLatch chunkStarted = new CountDownLatch(1); + SettableApiFuture> inFlightChunkFuture = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenAnswer( + inv -> { + chunkStarted.countDown(); + return inFlightChunkFuture; + }); + + TrackableStream stream = new TrackableStream("hello"); + ResumableUploadFuture future = callable.futureCall("resource-path", stream, null); + + // Wait until chunk callable is invoked so cancel tests in-flight chunk cancellation. + assertThat(chunkStarted.await(5, TimeUnit.SECONDS)).isTrue(); + + // While chunk upload is in-flight, cancel the session future + assertThat(future.cancel(true)).isTrue(); + assertThat(future.isCancelled()).isTrue(); + assertThat(future.isDone()).isTrue(); + + // The in-flight HTTP chunk future must have been cancelled + assertThat(inFlightChunkFuture.isCancelled()).isTrue(); + // The payload stream must be closed + assertThat(stream.closed).isTrue(); + } + + @Test + void testChunkRetry_cancellationDuringBackoff_deschedulesPendingAttempt() { + SettableApiFuture> chunkAttempt0Future = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(chunkAttempt0Future) + .thenReturn( + ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "should-not-reach"))); + + ClientContext clientContext = + ClientContext.newBuilder().setDefaultCallContext(callContext).build(); + // Leave startFuture incomplete so sessionFuture does not automatically instantiate its + // coordinator + SettableApiFuture startFuture = SettableApiFuture.create(); + + SettableApiFuture result = SettableApiFuture.create(); + SettableApiFuture startSessionFuture = SettableApiFuture.create(); + startSessionFuture.set( + ResumableUploadSession.newBuilder() + .setUploadUrl("https://upload.url/cancel-backoff") + .build()); + ResumableUploadCallSettings customSettings = + defaultSettings.toBuilder().setChunkSize(8).build(); + ResumableUploadChunkCoordinator coordinator = + new ResumableUploadChunkCoordinator<>( + result, + startSessionFuture, + mockChunkCallable, + streamOf("hello"), + customSettings, + callContext, + clientContext); + ResumableUploadFutureImpl sessionFuture = + new ResumableUploadFutureImpl<>(result, coordinator); + + coordinator.start(); + + // 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); + } + + @Test + void testTransmitChunk_largePayloadSmallChunksSynchronousCompletion_doesNotStackOverflow() + throws Exception { + stubStartSession("https://upload.url/trampoline-test"); + byte[] payload = new byte[200 * 1024]; // 200 KB + int chunkSize = 16; // 12,800 chunks + + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenAnswer( + inv -> { + ChunkUploadRequest req = inv.getArgument(0); + boolean isComplete = req.isFinal(); + return ApiFutures.immediateFuture( + ChunkUploadResponse.create(isComplete, isComplete ? "done-200k" : null)); + }); + + ResumableUploadCallSettings settings = + ResumableUploadCallSettings.newBuilder().setChunkSize(chunkSize).build(); + ResumableUploadFuture future = + callable.futureCall("resource-path", new ByteArrayInputStream(payload), settings); + + assertThat(future.get(10, TimeUnit.SECONDS)).isEqualTo("done-200k"); + } + private void stubStartSession(String uploadUrl) { when(mockStartCallable.futureCall(any(), any())) .thenReturn( @@ -378,4 +597,21 @@ public void close() throws IOException { super.close(); } } + + private static class ByteCountingStream extends ByteArrayInputStream { + int totalBytesRead = 0; + + ByteCountingStream(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; + } + } } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java new file mode 100644 index 000000000000..bdc9c8f36f80 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinatorTest.java @@ -0,0 +1,146 @@ +/* + * Copyright 2026 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.rpc; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +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.ResumableUploadSession; +import com.google.api.gax.rpc.testing.FakeCallContext; +import com.google.common.util.concurrent.MoreExecutors; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class ResumableUploadChunkCoordinatorTest { + + private ScheduledExecutorService executor; + private ClientContext clientContext; + private ApiCallContext callContext; + private UnaryCallable> mockChunkCallable; + private ResumableUploadCallSettings settings; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + executor = Executors.newScheduledThreadPool(2); + callContext = FakeCallContext.createDefault(); + clientContext = + ClientContext.newBuilder().setDefaultCallContext(callContext).setExecutor(executor).build(); + mockChunkCallable = mock(UnaryCallable.class, withSettings().withoutAnnotations()); + settings = ResumableUploadCallSettings.newBuilder().setChunkSize(256).build(); + } + + @AfterEach + void tearDown() { + executor.shutdownNow(); + } + + @Test + void testTerminalArbitration_concurrentCancelAndComplete_happensExactlyOnce() throws Exception { + SettableApiFuture result = SettableApiFuture.create(); + SettableApiFuture startFuture = SettableApiFuture.create(); + startFuture.set( + ResumableUploadSession.newBuilder() + .setUploadUrl("https://upload.url/arbitration-test") + .build()); + + SettableApiFuture> chunkFuture = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(chunkFuture); + + AtomicInteger closeCount = new AtomicInteger(0); + ByteArrayInputStream payload = + new ByteArrayInputStream(new byte[128]) { + @Override + public void close() throws IOException { + closeCount.incrementAndGet(); + super.close(); + } + }; + + ResumableUploadChunkCoordinator coordinator = + new ResumableUploadChunkCoordinator<>( + result, startFuture, mockChunkCallable, payload, settings, callContext, clientContext); + + AtomicInteger completionListenerCount = new AtomicInteger(0); + result.addListener(completionListenerCount::incrementAndGet, MoreExecutors.directExecutor()); + + coordinator.start(); + + CountDownLatch startLatch = new CountDownLatch(1); + CountDownLatch doneLatch = new CountDownLatch(2); + + executor.execute( + () -> { + try { + startLatch.await(); + coordinator.cancel(true); + result.cancel(true); + } catch (Exception e) { + // ignore + } finally { + doneLatch.countDown(); + } + }); + + executor.execute( + () -> { + try { + startLatch.await(); + chunkFuture.set(ChunkUploadResponse.create(true, "done", "final")); + } catch (Exception e) { + // ignore + } finally { + doneLatch.countDown(); + } + }); + + startLatch.countDown(); + assertThat(doneLatch.await(5, TimeUnit.SECONDS)).isTrue(); + + assertThat(result.isDone()).isTrue(); + assertThat(closeCount.get()).isEqualTo(1); + assertThat(completionListenerCount.get()).isEqualTo(1); + } +} diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadFutureImplTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadFutureImplTest.java new file mode 100644 index 000000000000..f5d7726a3d5c --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadFutureImplTest.java @@ -0,0 +1,122 @@ +/* + * Copyright 2026 Google LLC + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are + * met: + * + * * Redistributions of source code must retain the above copyright + * notice, this list of conditions and the following disclaimer. + * * Redistributions in binary form must reproduce the above + * copyright notice, this list of conditions and the following disclaimer + * in the documentation and/or other materials provided with the + * distribution. + * * Neither the name of Google LLC nor the names of its + * contributors may be used to endorse or promote products derived from + * this software without specific prior written permission. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + */ +package com.google.api.gax.rpc; + +import static com.google.common.truth.Truth.assertThat; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +import com.google.api.core.ApiFuture; +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.ResumableUploadSession; +import com.google.api.gax.rpc.testing.FakeCallContext; +import com.google.common.util.concurrent.MoreExecutors; +import java.io.ByteArrayInputStream; +import java.util.concurrent.atomic.AtomicBoolean; +import org.junit.jupiter.api.Test; + +class ResumableUploadFutureImplTest { + + @SuppressWarnings("unchecked") + private ResumableUploadChunkCoordinator createCoordinator( + SettableApiFuture result, + ApiFuture startFuture, + SettableApiFuture> chunkFuture) { + ClientContext clientContext = + ClientContext.newBuilder().setDefaultCallContext(FakeCallContext.createDefault()).build(); + UnaryCallable> mockChunkCallable = + mock(UnaryCallable.class, withSettings().withoutAnnotations()); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(chunkFuture); + ResumableUploadCallSettings settings = ResumableUploadCallSettings.newBuilder().build(); + return new ResumableUploadChunkCoordinator<>( + result, + startFuture, + mockChunkCallable, + new ByteArrayInputStream(new byte[] {1, 2, 3}), + settings, + clientContext.getDefaultCallContext(), + clientContext); + } + + @Test + void handle_delegatesToResultAndCoordinator() throws Exception { + SettableApiFuture result = SettableApiFuture.create(); + SettableApiFuture startFuture = SettableApiFuture.create(); + SettableApiFuture> chunkFuture = SettableApiFuture.create(); + ResumableUploadChunkCoordinator coordinator = + createCoordinator(result, startFuture, chunkFuture); + + ResumableUploadFutureImpl handle = new ResumableUploadFutureImpl<>(result, coordinator); + + assertThat(handle.getUploadSessionUrl()).isNull(); + assertThat(handle.isDone()).isFalse(); + assertThat(handle.isCancelled()).isFalse(); + + coordinator.start(); + startFuture.set( + ResumableUploadSession.newBuilder() + .setUploadUrl("https://upload.url/test-session") + .build()); + assertThat(handle.getUploadSessionUrl()).isEqualTo("https://upload.url/test-session"); + + AtomicBoolean listenerFired = new AtomicBoolean(false); + handle.addListener(() -> listenerFired.set(true), MoreExecutors.directExecutor()); + + result.set("response-ok"); + + assertThat(handle.isDone()).isTrue(); + assertThat(handle.get()).isEqualTo("response-ok"); + assertThat(listenerFired.get()).isTrue(); + } + + @Test + void handle_cancel_delegatesToCoordinatorAndResult() { + SettableApiFuture result = SettableApiFuture.create(); + SettableApiFuture startFuture = SettableApiFuture.create(); + SettableApiFuture> chunkFuture = SettableApiFuture.create(); + ResumableUploadChunkCoordinator coordinator = + createCoordinator(result, startFuture, chunkFuture); + coordinator.start(); + + ResumableUploadFutureImpl handle = new ResumableUploadFutureImpl<>(result, coordinator); + + boolean cancelled = handle.cancel(true); + + assertThat(cancelled).isTrue(); + assertThat(handle.isCancelled()).isTrue(); + assertThat(result.isCancelled()).isTrue(); + assertThat(startFuture.isCancelled()).isTrue(); + } +}