feat(gax): implement baseline Callable and Future for resumable uploads #14241
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 <RequestT> the type of the initial request message that initiates the upload session | ||
| * @param <ResponseT> the type of the final response message returned once the upload completes | ||
| */ | ||
| @BetaApi | ||
| @InternalApi | ||
| @NullMarked | ||
| public class ResumableUploadCallableImpl<RequestT, ResponseT> | ||
|
Check warning on line 55 in sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java
|
||
| extends ResumableUploadCallable<RequestT, ResponseT> { | ||
|
|
||
| private final ResumableUploadClient<RequestT, ResponseT> client; | ||
| private final ResumableUploadCallSettings defaultCallSettings; | ||
| private final ApiCallContext defaultCallContext; | ||
| private final ScheduledExecutorService executor; | ||
|
|
||
| public ResumableUploadCallableImpl( | ||
| ResumableUploadClient<RequestT, ResponseT> 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<ResponseT> 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<ResumableUploadSession> startFuture; | ||
| try { | ||
| startFuture = client.startUploadCallable().futureCall(request, defaultCallContext); | ||
| } catch (Throwable t) { | ||
|
Check warning on line 86 in sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadCallableImpl.java
|
||
| startFuture = ApiFutures.immediateFailedFuture(t); | ||
| } | ||
|
|
||
| return ResumableUploadFutureImpl.create( | ||
| startFuture, | ||
| client.uploadChunkCallable(), | ||
| payload, | ||
| effectiveSettings, | ||
| defaultCallContext, | ||
| executor); | ||
| } | ||
|
|
||
| @Override | ||
| public ResumableUploadFuture<ResponseT> resumeCall( | ||
| String sessionUrl, InputStream payload, @Nullable ResumableUploadCallSettings settings) { | ||
| throw new UnsupportedOperationException("Session resumption is not yet implemented."); | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 <ResponseT> the type of the final response message returned once the upload completes | ||
| */ | ||
| @InternalApi | ||
| @NullMarked | ||
| final class ResumableUploadChunkCoordinator<ResponseT> { | ||
|
Check warning on line 55 in sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java
|
||
|
|
||
| private static final byte[] EMPTY_PAYLOAD = new byte[0]; | ||
|
|
||
| private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> | ||
| uploadChunkCallable; | ||
| private final String uploadUrl; | ||
| private final InputStream payload; | ||
| private final int chunkSize; | ||
| private final ApiCallContext callContext; | ||
| private final ScheduledExecutorService executor; | ||
| private final ResumableUploadFutureImpl<ResponseT> sessionFuture; | ||
|
|
||
| ResumableUploadChunkCoordinator( | ||
| UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable, | ||
| String uploadUrl, | ||
| InputStream payload, | ||
| int chunkSize, | ||
| ApiCallContext callContext, | ||
| ScheduledExecutorService executor, | ||
| ResumableUploadFutureImpl<ResponseT> 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); | ||
|
Contributor
There was a problem hiding this comment. IIUC, this copy is to prevent corrupted stream at the end? Can we do it without copying? For example, using |
||
| } | ||
|
|
||
| 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<ChunkUploadResponse<ResponseT>> chunkFuture = | ||
| uploadChunkCallable.futureCall(chunkRequest, callContext); | ||
| sessionFuture.setInFlightFuture(chunkFuture); | ||
|
|
||
| // Asynchronously handle the response: complete, fail, or chain the next chunk. | ||
| ApiFutures.addCallback( | ||
| chunkFuture, | ||
| new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() { | ||
| @Override | ||
| public void onSuccess(ChunkUploadResponse<ResponseT> 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); | ||
|
Contributor
There was a problem hiding this comment. IIUC, the executor passed in is actually the background executor. It has a small number of thread pool and is mainly used for scheduling retries. If we use it for heavy operations like reading a stream, it would cause thread starvation pretty soon. Can we use In addition, I don't think we need to pass executor in from ClientContext unless we think it could be useful for scheduling retries later for recovering chunk uploading. |
||
| } catch (Throwable t) { | ||
|
Check warning on line 163 in sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java
|
||
| sessionFuture.fail(t); | ||
| } | ||
| } | ||
| } | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This new array allocation would be called for every chunk, which could cause high GC pressure. Can we reuse the same
byte[]for every chunk?