diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallable.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallable.java index 1aefa4adc6e5..8d4e441853dd 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallable.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallable.java @@ -31,15 +31,18 @@ import com.google.api.core.BetaApi; import java.io.InputStream; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /** * A ResumableUploadCallable is an API-transport-independent wrapper for the Resumable Upload * protocol. Operates directly on the request object and input stream payload. * - * @param request type - * @param response type + * @param the type of the initial request message that initiates the upload session + * @param the type of the final response message returned once the upload completes */ @BetaApi +@NullMarked public abstract class ResumableUploadCallable { protected ResumableUploadCallable() {} @@ -47,22 +50,30 @@ protected ResumableUploadCallable() {} /** * Performs a new resumable upload asynchronously. * + *

The provided {@code payload} stream is consumed asynchronously by the returned {@link + * ResumableUploadFuture} and will be closed automatically upon completion, failure, or + * cancellation. + * * @param request the request message - * @param payload the data payload input stream + * @param payload the data payload input stream to upload and close * @param settings call settings overrides; may be {@code null} * @return future for tracking and controlling the upload */ public abstract ResumableUploadFuture futureCall( - RequestT request, InputStream payload, ResumableUploadCallSettings settings); + RequestT request, InputStream payload, @Nullable ResumableUploadCallSettings settings); /** * Resumes an existing resumable upload session asynchronously using a saved session URL. * + *

The provided {@code payload} stream is consumed asynchronously by the returned {@link + * ResumableUploadFuture} and will be closed automatically upon completion, failure, or + * cancellation. + * * @param sessionUrl the upload session URL - * @param payload the data payload input stream + * @param payload the data payload input stream to upload and close * @param settings call settings overrides; may be {@code null} * @return future for tracking and controlling the upload */ public abstract ResumableUploadFuture resumeCall( - String sessionUrl, InputStream payload, ResumableUploadCallSettings settings); + String sessionUrl, InputStream payload, @Nullable ResumableUploadCallSettings settings); } 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 new file mode 100644 index 000000000000..aca5d1684287 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java @@ -0,0 +1,104 @@ +/* + * 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.base.Preconditions.checkNotNull; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutures; +import com.google.api.core.BetaApi; +import com.google.api.core.InternalApi; +import com.google.api.gax.resumable.ResumableUploadClient; +import com.google.api.gax.resumable.ResumableUploadSession; +import java.io.InputStream; +import java.util.concurrent.ScheduledExecutorService; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; + +/** + * Concrete implementation of {@link ResumableUploadCallable} that delegates the end-to-end + * management of a resumable upload session to {@link ResumableUploadFutureImpl}. + * + * @param the type of the initial request message that initiates the upload session + * @param the type of the final response message returned once the upload completes + */ +@BetaApi +@InternalApi +@NullMarked +public class ResumableUploadCallableImpl + extends ResumableUploadCallable { + + private final ResumableUploadClient client; + private final ResumableUploadCallSettings defaultCallSettings; + private final ApiCallContext defaultCallContext; + private final ScheduledExecutorService executor; + + public ResumableUploadCallableImpl( + ResumableUploadClient client, + ResumableUploadCallSettings defaultCallSettings, + ApiCallContext defaultCallContext, + ScheduledExecutorService executor) { + this.client = checkNotNull(client, "client must not be null"); + this.defaultCallSettings = + checkNotNull(defaultCallSettings, "defaultCallSettings must not be null"); + this.defaultCallContext = + checkNotNull(defaultCallContext, "defaultCallContext must not be null"); + this.executor = checkNotNull(executor, "executor must not be null"); + } + + @Override + public ResumableUploadFuture futureCall( + RequestT request, InputStream payload, @Nullable ResumableUploadCallSettings settings) { + checkNotNull(request, "request must not be null"); + checkNotNull(payload, "payload must not be null"); + ResumableUploadCallSettings effectiveSettings = defaultCallSettings.merge(settings); + + ApiFuture startFuture; + try { + startFuture = client.startUploadCallable().futureCall(request, defaultCallContext); + } catch (Throwable t) { + startFuture = ApiFutures.immediateFailedFuture(t); + } + + return ResumableUploadFutureImpl.create( + startFuture, + client.uploadChunkCallable(), + payload, + effectiveSettings, + defaultCallContext, + executor); + } + + @Override + public ResumableUploadFuture resumeCall( + String sessionUrl, InputStream payload, @Nullable ResumableUploadCallSettings settings) { + throw new UnsupportedOperationException("Session resumption is not yet implemented."); + } +} 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 new file mode 100644 index 000000000000..3a8129926ee4 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java @@ -0,0 +1,167 @@ +/* + * 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.base.Preconditions.checkNotNull; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutureCallback; +import com.google.api.core.ApiFutures; +import com.google.api.core.InternalApi; +import com.google.api.gax.resumable.ChunkUploadRequest; +import com.google.api.gax.resumable.ChunkUploadResponse; +import com.google.common.io.ByteStreams; +import java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ScheduledExecutorService; +import org.jspecify.annotations.NullMarked; + +/** + * Coordinates chunk transmission steps of a resumable upload session. + * + * @param the type of the final response message returned once the upload completes + */ +@InternalApi +@NullMarked +final class ResumableUploadChunkCoordinator { + + private static final byte[] EMPTY_PAYLOAD = new byte[0]; + + private final UnaryCallable> + uploadChunkCallable; + private final String uploadUrl; + private final InputStream payload; + private final int chunkSize; + private final ApiCallContext callContext; + private final ScheduledExecutorService executor; + private final ResumableUploadFutureImpl sessionFuture; + + ResumableUploadChunkCoordinator( + UnaryCallable> uploadChunkCallable, + String uploadUrl, + InputStream payload, + int chunkSize, + ApiCallContext callContext, + ScheduledExecutorService executor, + ResumableUploadFutureImpl sessionFuture) { + this.uploadChunkCallable = + checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); + this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null"); + this.payload = checkNotNull(payload, "payload must not be null"); + this.chunkSize = chunkSize; + this.callContext = checkNotNull(callContext, "callContext must not be null"); + this.executor = checkNotNull(executor, "executor must not be null"); + this.sessionFuture = checkNotNull(sessionFuture, "sessionFuture must not be null"); + } + + void start() { + transmitChunk(0L); + } + + private void transmitChunk(long currentOffset) { + // Abort if the session was already completed or canceled. + if (sessionFuture.isDone()) { + return; + } + + // Read the next chunk slice from the payload stream. + byte[] buffer = new byte[chunkSize]; + int bytesRead; + try { + bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize); + } catch (IOException e) { + sessionFuture.fail(e); + return; + } + + // Determine if this is the final chunk and build the chunk request. + boolean isFinal = bytesRead < chunkSize; + byte[] chunkPayload; + if (bytesRead == chunkSize) { + chunkPayload = buffer; + } else if (bytesRead == 0) { + chunkPayload = EMPTY_PAYLOAD; + } else { + chunkPayload = Arrays.copyOf(buffer, bytesRead); + } + + ChunkUploadRequest chunkRequest = + ChunkUploadRequest.newBuilder() + .setUploadUrl(uploadUrl) + .setPayload(chunkPayload) + .setOffset(currentOffset) + .setFinal(isFinal) + .build(); + + // Dispatch the chunk upload call and register the in-flight future for cancellation. + long chunkLength = chunkPayload.length; + try { + ApiFuture> chunkFuture = + uploadChunkCallable.futureCall(chunkRequest, callContext); + sessionFuture.setInFlightFuture(chunkFuture); + + // Asynchronously handle the response: complete, fail, or chain the next chunk. + ApiFutures.addCallback( + chunkFuture, + new ApiFutureCallback>() { + @Override + public void onSuccess(ChunkUploadResponse response) { + if (sessionFuture.isDone()) { + return; + } + long nextOffset = currentOffset + chunkLength; + if (response.isComplete()) { + sessionFuture.succeed(response.getResponse()); + } else if (isFinal) { + sessionFuture.fail( + new IllegalStateException( + "Upload stream ended and final chunk was transmitted, but server returned" + + " incomplete status")); + } else { + transmitChunk(nextOffset); + } + } + + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException || sessionFuture.isDone()) { + return; + } + sessionFuture.fail(t); + } + }, + executor); + } catch (Throwable t) { + sessionFuture.fail(t); + } + } +} diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java index ea2605e3a5d7..4282ef89fa33 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java @@ -31,15 +31,21 @@ import com.google.api.core.ApiFuture; import com.google.api.core.BetaApi; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /** * A specialized {@link ApiFuture} for tracking and controlling an in-flight resumable upload. * - * @param response type + *

The payload {@link java.io.InputStream} supplied when initiating the upload is managed by this + * future and will be closed automatically upon completion, failure, or cancellation. + * + * @param the type of the final response message returned once the upload completes */ @BetaApi +@NullMarked public interface ResumableUploadFuture extends ApiFuture { /** Returns the upload session URL, or {@code null} if session initiation is in progress. */ - String getUploadSessionUrl(); + @Nullable String getUploadSessionUrl(); } 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 new file mode 100644 index 000000000000..5d239b936a0d --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java @@ -0,0 +1,247 @@ +/* + * 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.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; + +import com.google.api.core.ApiFuture; +import com.google.api.core.ApiFutureCallback; +import com.google.api.core.ApiFutures; +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.errorprone.annotations.concurrent.GuardedBy; +import java.io.IOException; +import java.io.InputStream; +import java.util.concurrent.CancellationException; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executor; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; + +/** + * Implementation of {@link ResumableUploadFuture} responsible for the end-to-end management of a + * resumable upload session. + * + * @param the type of the final response message returned once the upload completes + */ +@NullMarked +final class ResumableUploadFutureImpl implements ResumableUploadFuture { + + private final Object lock = new Object(); + + private final ApiFuture startFuture; + private final UnaryCallable> + uploadChunkCallable; + 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; + + @GuardedBy("lock") + private @Nullable ApiFuture inFlightFuture; + + /** + * Creates and initiates a new resumable upload future tracking session initiation and chunk + * streaming. + * + *

The provided {@code payload} stream is managed by the returned future and will be closed + * automatically upon completion, failure, or cancellation. + */ + static ResumableUploadFutureImpl create( + ApiFuture startFuture, + UnaryCallable> uploadChunkCallable, + InputStream payload, + ResumableUploadCallSettings settings, + ApiCallContext callContext, + ScheduledExecutorService executor) { + ResumableUploadFutureImpl future = + new ResumableUploadFutureImpl<>( + startFuture, uploadChunkCallable, payload, settings, callContext, executor); + try { + future.start(); + } catch (Throwable t) { + future.fail(t); + } + return future; + } + + private ResumableUploadFutureImpl( + ApiFuture startFuture, + UnaryCallable> uploadChunkCallable, + InputStream payload, + ResumableUploadCallSettings settings, + ApiCallContext callContext, + ScheduledExecutorService executor) { + this.startFuture = checkNotNull(startFuture, "startFuture must not be null"); + this.uploadChunkCallable = + checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); + this.payload = checkNotNull(payload, "payload must not be null"); + this.settings = checkNotNull(settings, "settings must not be null"); + checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0"); + this.callContext = checkNotNull(callContext, "callContext must not be null"); + this.executor = checkNotNull(executor, "executor must not be null"); + this.inFlightFuture = startFuture; + } + + private void start() { + ApiFutures.addCallback( + startFuture, + new ApiFutureCallback() { + @Override + public void onSuccess(ResumableUploadSession session) { + if (resultFuture.isDone()) { + return; + } + uploadSessionUrl = session.getUploadUrl(); + ResumableUploadChunkCoordinator coordinator = + new ResumableUploadChunkCoordinator<>( + uploadChunkCallable, + session.getUploadUrl(), + payload, + settings.getChunkSize(), + callContext, + executor, + ResumableUploadFutureImpl.this); + try { + coordinator.start(); + } catch (Throwable t) { + fail(t); + } + } + + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException || resultFuture.isDone()) { + return; + } + fail(t); + } + }, + executor); + } + + /** + * Registers the active in-flight future for cancellation. If this session future has already been + * canceled, the supplied future is canceled immediately. + */ + void setInFlightFuture(ApiFuture inFlightFuture) { + boolean shouldCancel = false; + synchronized (lock) { + if (resultFuture.isDone()) { + shouldCancel = resultFuture.isCancelled(); + } else { + this.inFlightFuture = inFlightFuture; + } + } + if (shouldCancel) { + inFlightFuture.cancel(true); + } + } + + void succeed(@Nullable ResponseT result) { + synchronized (lock) { + inFlightFuture = null; + } + closePayload(); + resultFuture.set(result); + } + + void fail(Throwable t) { + synchronized (lock) { + inFlightFuture = null; + } + closePayload(); + resultFuture.setException(t); + } + + private void closePayload() { + try { + payload.close(); + } catch (IOException ignored) { + // Suppressed during stream cleanup + } + } + + @Override + public @Nullable String getUploadSessionUrl() { + return uploadSessionUrl; + } + + @Override + public void addListener(Runnable listener, Executor executor) { + resultFuture.addListener(listener, executor); + } + + @Override + public boolean cancel(boolean mayInterruptIfRunning) { + boolean cancelled; + ApiFuture inFlight; + synchronized (lock) { + cancelled = resultFuture.cancel(mayInterruptIfRunning); + inFlight = this.inFlightFuture; + this.inFlightFuture = null; + } + if (inFlight != null) { + inFlight.cancel(mayInterruptIfRunning); + } + closePayload(); + return cancelled; + } + + @Override + public boolean isCancelled() { + return resultFuture.isCancelled(); + } + + @Override + public boolean isDone() { + return resultFuture.isDone(); + } + + @Override + public ResponseT get() throws InterruptedException, ExecutionException { + return resultFuture.get(); + } + + @Override + public ResponseT get(long timeout, TimeUnit unit) + throws InterruptedException, ExecutionException, TimeoutException { + return resultFuture.get(timeout, unit); + } +} 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 new file mode 100644 index 000000000000..f683fcb638d0 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java @@ -0,0 +1,391 @@ +/* + * 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.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +import com.google.api.core.ApiFutures; +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.ResumableUploadClient; +import com.google.api.gax.resumable.ResumableUploadSession; +import com.google.api.gax.rpc.testing.FakeCallContext; +import java.io.ByteArrayInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import java.util.List; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledThreadPoolExecutor; +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; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class ResumableUploadCallableImplTest { + + private ResumableUploadClient mockClient; + private UnaryCallable mockStartCallable; + private UnaryCallable> mockChunkCallable; + + private ResumableUploadCallSettings defaultSettings; + private FakeCallContext callContext; + private ScheduledExecutorService executor; + private ResumableUploadCallableImpl callable; + + @BeforeEach + @SuppressWarnings("unchecked") + void setUp() { + mockClient = mock(ResumableUploadClient.class, withSettings().withoutAnnotations()); + mockStartCallable = mock(UnaryCallable.class, withSettings().withoutAnnotations()); + mockChunkCallable = mock(UnaryCallable.class, withSettings().withoutAnnotations()); + + lenient().when(mockClient.startUploadCallable()).thenReturn(mockStartCallable); + lenient().when(mockClient.uploadChunkCallable()).thenReturn(mockChunkCallable); + + defaultSettings = ResumableUploadCallSettings.newBuilder().setChunkSize(8).build(); + callContext = FakeCallContext.createDefault(); + executor = Executors.newSingleThreadScheduledExecutor(); + callable = + new ResumableUploadCallableImpl<>(mockClient, defaultSettings, callContext, executor); + } + + @AfterEach + void tearDown() throws InterruptedException { + if (executor != null) { + executor.shutdownNow(); + executor.awaitTermination(5, TimeUnit.SECONDS); + } + } + + @Test + void testUploadCallable_singleChunk_happyPath() throws Exception { + stubStartSession("https://upload.url/single"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "response-single"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("hello"), null); + + assertThat(future.get()).isEqualTo("response-single"); + assertThat(future.isDone()).isTrue(); + assertThat(future.isCancelled()).isFalse(); + assertThat(future.getUploadSessionUrl()).isEqualTo("https://upload.url/single"); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable).futureCall(captor.capture(), any()); + ChunkUploadRequest chunk = captor.getValue(); + assertThat(chunk.getUploadUrl()).isEqualTo("https://upload.url/single"); + assertChunk(chunk, 0, 5, true); + } + + @Test + void testUploadCallable_multiChunk_happyPath() throws Exception { + stubStartSession("https://upload.url/multi"); + // 20 bytes with chunkSize = 8 -> 8 + 8 + 4 bytes + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "response-multi"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("01234567890123456789"), null); + + assertThat(future.get()).isEqualTo("response-multi"); + assertThat(future.isDone()).isTrue(); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(3)).futureCall(captor.capture(), any()); + List chunks = captor.getAllValues(); + assertChunk(chunks.get(0), 0, 8, false); + assertChunk(chunks.get(1), 8, 8, false); + assertChunk(chunks.get(2), 16, 4, true); + } + + @Test + void testUploadCallable_zeroByteUpload_finalizesSuccessfully() throws Exception { + stubStartSession("https://upload.url/zero"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "response-zero"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", new ByteArrayInputStream(new byte[0]), null); + + assertThat(future.get()).isEqualTo("response-zero"); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable).futureCall(captor.capture(), any()); + assertChunk(captor.getValue(), 0, 0, true); + } + + @Test + void testUploadCallable_singleChunkWithSeparateZeroByteFinalize_completesSuccessfully() + throws Exception { + stubStartSession("https://upload.url/exact-single"); + // Exactly 8 bytes with chunkSize = 8 -> 8 bytes (non-final) then 0 bytes (final) + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(false, null))) + .thenReturn( + ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "response-exact-single"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("12345678"), null); + + assertThat(future.get()).isEqualTo("response-exact-single"); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any()); + List chunks = captor.getAllValues(); + assertChunk(chunks.get(0), 0, 8, false); + assertChunk(chunks.get(1), 8, 0, true); + } + + @Test + void testUploadCallable_nullResponse_completesSuccessfully() throws Exception { + stubStartSession("https://upload.url/null-response"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, null))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("data"), null); + + assertThat(future.get()).isNull(); + assertThat(future.isDone()).isTrue(); + assertThat(future.isCancelled()).isFalse(); + } + + @Test + void testUploadCallable_cancelInFlight_haltsUpload() throws Exception { + stubStartSession("https://upload.url/cancel"); + + CountDownLatch chunkStarted = new CountDownLatch(1); + CountDownLatch chunkCancelled = new CountDownLatch(1); + SettableApiFuture> pendingChunkFuture = SettableApiFuture.create(); + pendingChunkFuture.addListener(chunkCancelled::countDown, Runnable::run); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenAnswer( + inv -> { + chunkStarted.countDown(); + return pendingChunkFuture; + }); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("0123456789ABCDEF"), null); + + assertThat(chunkStarted.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(future.cancel(true)).isTrue(); + assertThat(future.isCancelled()).isTrue(); + assertThat(future.isDone()).isTrue(); + assertThat(chunkCancelled.await(5, TimeUnit.SECONDS)).isTrue(); + assertThat(pendingChunkFuture.isCancelled()).isTrue(); + assertThrows(CancellationException.class, future::get); + } + + @Test + void testUploadCallable_setInFlightFutureAfterCancel_immediatelyCancelsFuture() { + SettableApiFuture startFuture = SettableApiFuture.create(); + when(mockStartCallable.futureCall(any(), any())).thenReturn(startFuture); + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("data"), null); + assertThat(future.cancel(true)).isTrue(); + assertThat(future.isCancelled()).isTrue(); + + SettableApiFuture lateFuture = SettableApiFuture.create(); + ((ResumableUploadFutureImpl) future).setInFlightFuture(lateFuture); + assertThat(lateFuture.isCancelled()).isTrue(); + } + + @Test + void testUploadCallable_startFailure_failsFuture() { + when(mockStartCallable.futureCall(any(), any())) + .thenReturn(ApiFutures.immediateFailedFuture(new IllegalStateException("start failed"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("data"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(exception.getCause()).hasMessageThat().contains("start failed"); + verifyNoInteractions(mockChunkCallable); + } + + @Test + void testUploadCallable_chunkFailure_failsFuture() { + stubStartSession("https://upload.url/chunk-fail"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFailedFuture(new IllegalStateException("chunk error"))); + + ResumableUploadFuture future = + callable.futureCall("resource-path", streamOf("data"), null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(IllegalStateException.class); + assertThat(exception.getCause()).hasMessageThat().contains("chunk error"); + } + + @Test + void testUploadCallable_closesPayloadOnSuccess() throws Exception { + stubStartSession("https://upload.url/close-success"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn(ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "done"))); + + TrackableStream stream = new TrackableStream("data"); + callable.futureCall("resource-path", stream, null).get(); + + assertThat(stream.closed).isTrue(); + } + + @Test + void testUploadCallable_closesPayloadOnFailure() { + when(mockStartCallable.futureCall(any(), any())) + .thenReturn(ApiFutures.immediateFailedFuture(new IllegalStateException("start failed"))); + + TrackableStream stream = new TrackableStream("data"); + ResumableUploadFuture future = callable.futureCall("resource-path", stream, null); + assertThrows(ExecutionException.class, future::get); + + assertThat(stream.closed).isTrue(); + } + + @Test + void testUploadCallable_closesPayloadOnCancel() throws Exception { + stubStartSession("https://upload.url/close-cancel"); + CountDownLatch chunkStarted = new CountDownLatch(1); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenAnswer( + inv -> { + chunkStarted.countDown(); + return SettableApiFuture.create(); + }); + + TrackableStream stream = new TrackableStream("data"); + ResumableUploadFuture future = callable.futureCall("resource-path", stream, null); + assertThat(chunkStarted.await(5, TimeUnit.SECONDS)).isTrue(); + future.cancel(true); + + assertThat(stream.closed).isTrue(); + } + + @Test + void testUploadCallable_closesPayloadOnStartSyncFailure() { + when(mockStartCallable.futureCall(any(), any())) + .thenThrow(new RuntimeException("sync start failure")); + + TrackableStream stream = new TrackableStream("data"); + ResumableUploadFuture future = callable.futureCall("resource-path", stream, null); + + ExecutionException exception = assertThrows(ExecutionException.class, future::get); + assertThat(exception.getCause()).isInstanceOf(RuntimeException.class); + assertThat(exception.getCause()).hasMessageThat().contains("sync start failure"); + assertThat(stream.closed).isTrue(); + } + + @Test + void testUploadCallable_customExecutor_executesOnExecutor() throws Exception { + AtomicInteger tasksRun = new AtomicInteger(); + ScheduledExecutorService customExecutor = + new ScheduledThreadPoolExecutor(1) { + @Override + public void execute(Runnable command) { + tasksRun.incrementAndGet(); + super.execute(command); + } + }; + try { + ResumableUploadCallableImpl customCallable = + new ResumableUploadCallableImpl<>( + mockClient, defaultSettings, callContext, customExecutor); + stubStartSession("https://upload.url/executor"); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture(ChunkUploadResponse.create(true, "executor-done"))); + + ResumableUploadFuture future = + customCallable.futureCall("resource-path", streamOf("data"), null); + assertThat(future.get()).isEqualTo("executor-done"); + assertThat(tasksRun.get()).isGreaterThan(0); + } finally { + customExecutor.shutdownNow(); + customExecutor.awaitTermination(5, TimeUnit.SECONDS); + } + } + + private void stubStartSession(String uploadUrl) { + when(mockStartCallable.futureCall(any(), any())) + .thenReturn( + ApiFutures.immediateFuture( + ResumableUploadSession.newBuilder().setUploadUrl(uploadUrl).build())); + } + + private static InputStream streamOf(String content) { + return new ByteArrayInputStream(content.getBytes(StandardCharsets.UTF_8)); + } + + private static void assertChunk( + ChunkUploadRequest chunk, long expectedOffset, int expectedSize, boolean expectedFinal) { + assertThat(chunk.getOffset()).isEqualTo(expectedOffset); + assertThat(chunk.getPayload().length).isEqualTo(expectedSize); + assertThat(chunk.isFinal()).isEqualTo(expectedFinal); + } + + private static class TrackableStream extends ByteArrayInputStream { + boolean closed = false; + + TrackableStream(String content) { + super(content.getBytes(StandardCharsets.UTF_8)); + } + + @Override + public void close() throws IOException { + closed = true; + super.close(); + } + } +}