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 b52912ba8ebb..42d10e56e695 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 @@ -119,7 +119,7 @@ public ResumableUploadFuture futureCall( retryingQueryCallable, payload, effectiveSettings, - clientContext.getDefaultCallContext()); + clientContext); } @Override 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 aa0ae5fe4bac..0d046edff2c7 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 @@ -29,6 +29,7 @@ */ package com.google.api.gax.rpc; +import static com.google.common.base.MoreObjects.firstNonNull; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; @@ -45,9 +46,12 @@ import com.google.errorprone.annotations.concurrent.GuardedBy; import java.io.IOException; import java.io.InputStream; +import java.time.Duration; import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import org.jspecify.annotations.NullMarked; @@ -62,6 +66,8 @@ @NullMarked final class ResumableUploadFutureImpl implements ResumableUploadFuture { + private static final Duration DEFAULT_GLOBAL_TIMEOUT = Duration.ofMinutes(15); + private final Object lock = new Object(); private final ApiFuture startFuture; @@ -72,6 +78,7 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur private final InputStream payload; private final ResumableUploadCallSettings settings; private final ApiCallContext callContext; + private final ScheduledExecutorService executor; private final SettableApiFuture resultFuture = SettableApiFuture.create(); private volatile @Nullable String uploadSessionUrl; @@ -92,10 +99,15 @@ static ResumableUploadFutureImpl create( UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, - ApiCallContext callContext) { + ClientContext clientContext) { ResumableUploadFutureImpl future = new ResumableUploadFutureImpl<>( - startFuture, uploadChunkCallable, queryStatusCallable, payload, settings, callContext); + startFuture, + uploadChunkCallable, + queryStatusCallable, + payload, + settings, + clientContext); try { future.start(); } catch (Throwable t) { @@ -110,7 +122,7 @@ private ResumableUploadFutureImpl( UnaryCallable> queryStatusCallable, InputStream payload, ResumableUploadCallSettings settings, - ApiCallContext callContext) { + ClientContext clientContext) { this.startFuture = checkNotNull(startFuture, "startFuture must not be null"); this.uploadChunkCallable = checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); @@ -119,11 +131,17 @@ private ResumableUploadFutureImpl( 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.inFlightFuture = startFuture; } private void start() { + Duration timeout = firstNonNull(settings.getGlobalTimeout(), DEFAULT_GLOBAL_TIMEOUT); + ScheduledFuture timeoutFuture = + executor.schedule(this::onTimeout, timeout.toMillis(), TimeUnit.MILLISECONDS); + resultFuture.addListener(() -> timeoutFuture.cancel(false), MoreExecutors.directExecutor()); ApiFutures.addCallback( startFuture, new ApiFutureCallback() { @@ -190,6 +208,17 @@ public void onFailure(Throwable t) { MoreExecutors.directExecutor()); } + private void onTimeout() { + String sessionUrl = uploadSessionUrl; + String message; + if (sessionUrl != null) { + message = "Resumable upload timed out for session: " + sessionUrl; + } else { + message = "Resumable upload timed out before session initiation completed"; + } + fail(new DeadlineExceededException(message, null, TIMEOUT_STATUS_CODE, false)); + } + private void succeed(@Nullable ResponseT result) { synchronized (lock) { inFlightFuture = null; @@ -199,8 +228,13 @@ private void succeed(@Nullable ResponseT result) { } private void fail(Throwable t) { + ApiFuture inFlight; synchronized (lock) { - inFlightFuture = null; + inFlight = this.inFlightFuture; + this.inFlightFuture = null; + } + if (inFlight != null) { + inFlight.cancel(true); } closePayload(); resultFuture.setException(t); @@ -260,4 +294,17 @@ public ResponseT get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { return resultFuture.get(timeout, unit); } + + private static final StatusCode TIMEOUT_STATUS_CODE = + new StatusCode() { + @Override + public StatusCode.Code getCode() { + return StatusCode.Code.DEADLINE_EXCEEDED; + } + + @Override + public @Nullable Object getTransportCode() { + return null; + } + }; } 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 aa6eb56e0237..c0f6f1b0f4ba 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 @@ -32,6 +32,8 @@ import static com.google.common.truth.Truth.assertThat; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; @@ -54,11 +56,14 @@ import java.io.IOException; import java.io.InputStream; import java.nio.charset.StandardCharsets; +import java.time.Duration; import java.util.Arrays; import java.util.List; import java.util.concurrent.CancellationException; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import org.jspecify.annotations.Nullable; import org.junit.jupiter.api.BeforeEach; @@ -77,6 +82,7 @@ class ResumableUploadCallableImplTest { private ResumableUploadCallSettings defaultSettings; private FakeCallContext callContext; + private ClientContext clientContext; private ResumableUploadCallableImpl callable; @BeforeEach @@ -93,8 +99,7 @@ void setUp() { defaultSettings = ResumableUploadCallSettings.newBuilder().setChunkSize(8).build(); callContext = FakeCallContext.createDefault(); - ClientContext clientContext = - ClientContext.newBuilder().setDefaultCallContext(callContext).build(); + clientContext = ClientContext.newBuilder().setDefaultCallContext(callContext).build(); callable = new ResumableUploadCallableImpl<>(mockClient, defaultSettings, clientContext); } @@ -742,6 +747,210 @@ void testRecovery_queryTransientError_retriesAndSucceeds() throws Exception { verify(mockChunkCallable, times(2)).futureCall(any(), any()); } + @Test + void testGlobalTimeout_firesAndFailsSessionWithDeadlineExceeded() throws Exception { + stubStartSession("https://upload.url/timeout-fire"); + SettableApiFuture> hungChunk = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())).thenReturn(hungChunk); + + ResumableUploadCallSettings timeoutSettings = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofMillis(100)).build(); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), timeoutSettings); + + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertThat(exception.getCause()).isInstanceOf(DeadlineExceededException.class); + DeadlineExceededException cause = (DeadlineExceededException) exception.getCause(); + assertThat(cause.getStatusCode().getCode()).isEqualTo(StatusCode.Code.DEADLINE_EXCEEDED); + assertThat(cause.getMessage()).contains("https://upload.url/timeout-fire"); + assertThat(future.isDone()).isTrue(); + assertThat(future.isCancelled()).isFalse(); + } + + @Test + void testGlobalTimeout_cancelledCleanlyOnSuccess() throws Exception { + stubStartSession("https://upload.url/timeout-success"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "ok"))); + + ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class); + ScheduledFuture mockScheduledFuture = mock(ScheduledFuture.class); + when(mockExecutor.schedule(any(Runnable.class), anyLong(), any())) + .thenAnswer(inv -> mockScheduledFuture); + + ClientContext customClientContext = clientContext.toBuilder().setExecutor(mockExecutor).build(); + ResumableUploadCallableImpl customCallable = + new ResumableUploadCallableImpl<>(mockClient, defaultSettings, customClientContext); + + ResumableUploadCallSettings timeoutSettings = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofSeconds(60)).build(); + + ResumableUploadFuture future = + customCallable.futureCall("resource-path", streamOf("hello"), timeoutSettings); + + assertThat(future.get()).isEqualTo("ok"); + verify(mockScheduledFuture).cancel(false); + } + + @Test + void testGlobalTimeout_cancelledCleanlyOnFailure() throws Exception { + when(mockStartCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(401, StatusCode.Code.UNAUTHENTICATED))); + + ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class); + ScheduledFuture mockScheduledFuture = mock(ScheduledFuture.class); + when(mockExecutor.schedule(any(Runnable.class), anyLong(), any())) + .thenAnswer(inv -> mockScheduledFuture); + + ClientContext customClientContext = clientContext.toBuilder().setExecutor(mockExecutor).build(); + ResumableUploadCallableImpl customCallable = + new ResumableUploadCallableImpl<>(mockClient, defaultSettings, customClientContext); + + ResumableUploadCallSettings timeoutSettings = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofSeconds(60)).build(); + + ResumableUploadFuture future = + customCallable.futureCall("resource-path", streamOf("hello"), timeoutSettings); + + assertThrows(ExecutionException.class, future::get); + verify(mockScheduledFuture).cancel(false); + } + + @Test + void testGlobalTimeout_cancelledCleanlyOnUserCancel() throws Exception { + stubStartSession("https://upload.url/timeout-cancel"); + SettableApiFuture> hungChunk = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())).thenReturn(hungChunk); + + ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class); + ScheduledFuture mockScheduledFuture = mock(ScheduledFuture.class); + when(mockExecutor.schedule(any(Runnable.class), anyLong(), any())) + .thenAnswer(inv -> mockScheduledFuture); + + ClientContext customClientContext = clientContext.toBuilder().setExecutor(mockExecutor).build(); + ResumableUploadCallableImpl customCallable = + new ResumableUploadCallableImpl<>(mockClient, defaultSettings, customClientContext); + + ResumableUploadCallSettings timeoutSettings = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofSeconds(60)).build(); + + ResumableUploadFuture future = + customCallable.futureCall("resource-path", streamOf("hello"), timeoutSettings); + + assertThat(future.cancel(true)).isTrue(); + verify(mockScheduledFuture).cancel(false); + } + + @Test + void testGlobalTimeout_usesDefaultWhenUnset() throws Exception { + stubStartSession("https://upload.url/default-timeout"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "ok"))); + + ScheduledExecutorService mockExecutor = mock(ScheduledExecutorService.class); + ScheduledFuture mockScheduledFuture = mock(ScheduledFuture.class); + when(mockExecutor.schedule(any(Runnable.class), anyLong(), any())) + .thenAnswer(inv -> mockScheduledFuture); + + ClientContext customClientContext = clientContext.toBuilder().setExecutor(mockExecutor).build(); + + // 1. Unset on both stub and per-request -> falls back to GAX default (15m) + ResumableUploadCallableImpl unsetStubCallable = + new ResumableUploadCallableImpl<>(mockClient, defaultSettings, customClientContext); + assertThat(unsetStubCallable.futureCall("resource-path", streamOf("hello"), null).get()) + .isEqualTo("ok"); + verify(mockExecutor) + .schedule( + any(Runnable.class), eq(Duration.ofMinutes(15).toMillis()), eq(TimeUnit.MILLISECONDS)); + + // 2. Stub-level timeout (30m, e.g. from generator/client settings) + null per-request -> 30m + ResumableUploadCallSettings stubWith30m = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofMinutes(30)).build(); + ResumableUploadCallableImpl configuredStubCallable = + new ResumableUploadCallableImpl<>(mockClient, stubWith30m, customClientContext); + assertThat(configuredStubCallable.futureCall("resource-path", streamOf("hello"), null).get()) + .isEqualTo("ok"); + verify(mockExecutor) + .schedule( + any(Runnable.class), eq(Duration.ofMinutes(30).toMillis()), eq(TimeUnit.MILLISECONDS)); + + // 3. Stub-level timeout (30m) + per-request with only chunkSize set -> preserves 30m + ResumableUploadCallSettings perRequestChunkSizeOnly = + ResumableUploadCallSettings.newBuilder().setChunkSize(16).build(); + assertThat( + configuredStubCallable + .futureCall("resource-path", streamOf("hello"), perRequestChunkSizeOnly) + .get()) + .isEqualTo("ok"); + verify(mockExecutor, times(2)) + .schedule( + any(Runnable.class), eq(Duration.ofMinutes(30).toMillis()), eq(TimeUnit.MILLISECONDS)); + + // 4. Stub-level timeout (30m) + per-request globalTimeout (5m) -> per-request wins (5m) + ResumableUploadCallSettings perRequestWith5m = + ResumableUploadCallSettings.newBuilder().setGlobalTimeout(Duration.ofMinutes(5)).build(); + assertThat( + configuredStubCallable + .futureCall("resource-path", streamOf("hello"), perRequestWith5m) + .get()) + .isEqualTo("ok"); + verify(mockExecutor) + .schedule( + any(Runnable.class), eq(Duration.ofMinutes(5).toMillis()), eq(TimeUnit.MILLISECONDS)); + } + + @Test + void testGlobalTimeout_timeoutWhileAttemptInFlight_cancelsInFlightFutureAndDoesNotCorruptBuffer() + throws Exception { + stubStartSession("https://upload.url/in-flight-timeout"); + SettableApiFuture> inFlightFuture = SettableApiFuture.create(); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(inFlightFuture); + + ResumableUploadCallSettings timeoutSettings = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofMillis(80)).build(); + + TrackableStream stream = new TrackableStream("01234567890123456789"); + ResumableUploadFuture future = + callable.futureCall("resource-path", stream, timeoutSettings); + + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertThat(exception.getCause()).isInstanceOf(DeadlineExceededException.class); + // In-flight attempt future must be cancelled + assertThat(inFlightFuture.isCancelled()).isTrue(); + + // Stream should have been read only up to the first chunk (chunkSize = 8), not refilled or + // advanced + assertThat(stream.totalBytesRead).isEqualTo(8); + } + + @Test + void testGlobalTimeout_coversStartSessionTimeout() throws Exception { + SettableApiFuture hungStartFuture = SettableApiFuture.create(); + when(mockStartCallable.futureCall(any(), any())).thenReturn(hungStartFuture); + + ResumableUploadCallSettings timeoutSettings = + defaultSettings.toBuilder().setGlobalTimeout(Duration.ofMillis(80)).build(); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), timeoutSettings); + + ExecutionException exception = + assertThrows(ExecutionException.class, () -> future.get(5, TimeUnit.SECONDS)); + assertThat(exception.getCause()).isInstanceOf(DeadlineExceededException.class); + assertThat(exception.getCause().getMessage()).contains("before session initiation completed"); + assertThat(hungStartFuture.isCancelled()).isTrue(); + } + private static class HttpStatusStatusCode implements StatusCode { private final int httpStatus; private final StatusCode.Code code;