diff --git a/java-showcase/gapic-showcase/src/test/java/com/google/showcase/v1beta1/it/ITResumableUpload.java b/java-showcase/gapic-showcase/src/test/java/com/google/showcase/v1beta1/it/ITResumableUpload.java index 871374fbc565..f259b6895578 100644 --- a/java-showcase/gapic-showcase/src/test/java/com/google/showcase/v1beta1/it/ITResumableUpload.java +++ b/java-showcase/gapic-showcase/src/test/java/com/google/showcase/v1beta1/it/ITResumableUpload.java @@ -359,6 +359,62 @@ void testProgressListener(@TempDir Path tempDir) throws Exception { } } + @Test + void testResumeUpload(@TempDir Path tempDir) throws Exception { + int totalBytes = 600 * 1024; + Path file = createTempFile(tempDir, "it-resume.txt", totalBytes); + UploadMediaRequest request = UploadMediaRequest.newBuilder().setName("it-resume.txt").build(); + ApiCallContext failSecondChunkContext = faultInjectionContext(403, 1, SHOWCASE_CHUNK_SIZE); + + ResumableUploadFuture firstAttemptFuture = + client + .uploadMediaCallable() + .futureCall( + request, + () -> Files.newInputStream(file), + failSecondChunkContext, + DEFAULT_TEST_OPTIONS); + + assertThrows(ExecutionException.class, () -> firstAttemptFuture.get(15, TimeUnit.SECONDS)); + String sessionUrl = firstAttemptFuture.getUploadSessionUrl(); + assertThat(sessionUrl).isNotNull(); + + List resumeProgress = new CopyOnWriteArrayList<>(); + CountDownLatch finalizedLatch = new CountDownLatch(1); + ResumableUploadOptions resumeOptions = + DEFAULT_TEST_OPTIONS.toBuilder() + .setProgressListener( + status -> { + resumeProgress.add(status); + if ("FINALIZED".equals(status.getState())) { + finalizedLatch.countDown(); + } + }) + .build(); + + ResumableUploadFuture secondAttemptFuture = + client + .uploadMediaCallable() + .resumeCall(sessionUrl, () -> Files.newInputStream(file), resumeOptions); + UploadMediaResponse response = secondAttemptFuture.get(30, TimeUnit.SECONDS); + assertThat(finalizedLatch.await(10, TimeUnit.SECONDS)).isTrue(); + + assertThat(response.getName()).isEqualTo("it-resume.txt"); + assertThat(response.getSize()).isEqualTo(totalBytes); + assertThat( + resumeProgress.stream() + .map(ResumableUploadProgress::getState) + .collect(Collectors.toList())) + .containsExactly("RECOVERING", "OFFSET_RECEIVED", "UPLOADING", "FINALIZED") + .inOrder(); + assertThat( + resumeProgress.stream() + .map(ResumableUploadProgress::getBytesUploaded) + .collect(Collectors.toList())) + .containsExactly(0L, 262144L, 524288L, 614400L) + .inOrder(); + } + private static Path createTempFile(Path dir, String fileName, byte[] data) throws IOException { Path path = dir.resolve(fileName); Files.write(path, data); diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/InputStreamSupplier.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/InputStreamSupplier.java index ceb6a9b8a347..3a0f8d9311a8 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/InputStreamSupplier.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/InputStreamSupplier.java @@ -43,13 +43,14 @@ public interface InputStreamSupplier { /** - * Opens or returns the {@link InputStream} to be read. + * Opens or returns the {@link InputStream} to be read, positioned at the beginning of the + * payload. * *

The returned stream is managed by the resumable upload and is closed automatically when the * upload succeeds, fails, times out, or is cancelled. Callers should not close the stream * themselves. * - * @return a non-null input stream over the payload data + * @return a non-null input stream positioned at byte {@code 0} of the payload data * @throws IOException if an I/O error occurs when opening or obtaining the stream */ InputStream get() throws IOException; 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 1418d3fb3bbe..425314d78e7b 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 @@ -49,51 +49,59 @@ protected ResumableUploadCallable() {} /** * Performs a new resumable upload asynchronously with options overrides and default call context. * - *

Streams from the provided {@code payload} are consumed asynchronously by the returned {@link - * ResumableUploadFuture} and will be closed automatically upon completion, failure, or - * cancellation. + *

Streams from the provided {@code payloadSupplier} are consumed asynchronously by the + * returned {@link ResumableUploadFuture} and will be closed automatically upon completion, + * failure, or cancellation. * * @param request the request message - * @param payload supplier of the data payload input stream to upload and close + * @param payloadSupplier supplier of the data payload input stream to upload and close * @param options request-level upload options overrides; may be {@code null} * @return future for tracking and controlling the upload */ public ResumableUploadFuture futureCall( - RequestT request, InputStreamSupplier payload, @Nullable ResumableUploadOptions options) { - return futureCall(request, payload, null, options); + RequestT request, + InputStreamSupplier payloadSupplier, + @Nullable ResumableUploadOptions options) { + return futureCall(request, payloadSupplier, null, options); } /** * Performs a new resumable upload asynchronously with call context and options overrides. * - *

Streams from the provided {@code payload} are consumed asynchronously by the returned {@link - * ResumableUploadFuture} and will be closed automatically upon completion, failure, or - * cancellation. + *

Streams from the provided {@code payloadSupplier} are consumed asynchronously by the + * returned {@link ResumableUploadFuture} and will be closed automatically upon completion, + * failure, or cancellation. * * @param request the request message - * @param payload supplier of the data payload input stream to upload and close + * @param payloadSupplier supplier of the data payload input stream to upload and close * @param context call context overrides; may be {@code null} * @param options request-level upload options overrides; may be {@code null} * @return future for tracking and controlling the upload */ public abstract ResumableUploadFuture futureCall( RequestT request, - InputStreamSupplier payload, + InputStreamSupplier payloadSupplier, @Nullable ApiCallContext context, @Nullable ResumableUploadOptions options); /** * Resumes an existing resumable upload session asynchronously using a saved session URL. * - *

Streams from the provided {@code payload} are consumed asynchronously by the returned {@link - * ResumableUploadFuture} and will be closed automatically upon completion, failure, or - * cancellation. + *

Resuming queries the server for the last committed byte offset and skips the stream forward + * to that offset before resuming chunk transmission. {@code payloadSupplier} must supply a stream + * positioned at byte {@code 0} of the original payload. + * + *

Streams from the provided {@code payloadSupplier} are 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 supplier of the data payload input stream to upload and close + * @param payloadSupplier supplier of the data payload input stream to upload and close * @param options upload options overrides; may be {@code null} * @return future for tracking and controlling the upload */ public abstract ResumableUploadFuture resumeCall( - String sessionUrl, InputStreamSupplier payload, @Nullable ResumableUploadOptions options); + String sessionUrl, + InputStreamSupplier payloadSupplier, + @Nullable ResumableUploadOptions options); } 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 127678aa660f..2c9c2537a443 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 @@ -29,6 +29,7 @@ */ 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; @@ -45,6 +46,7 @@ import com.google.api.gax.retrying.RetryAlgorithm; import com.google.api.gax.retrying.RetrySettings; import com.google.api.gax.retrying.ScheduledRetryingExecutor; +import com.google.common.base.Strings; import java.time.Duration; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -62,7 +64,7 @@ public class ResumableUploadCallableImpl extends ResumableUploadCallable { - private static final RetrySettings RETRY_SETTINGS = + static final RetrySettings RETRY_SETTINGS = RetrySettings.newBuilder() .setInitialRetryDelayDuration(Duration.ofSeconds(1)) .setRetryDelayMultiplier(2.0) @@ -81,7 +83,6 @@ public class ResumableUploadCallableImpl retryingUploadChunkCallable; private final UnaryCallable> retryingQueryCallable; - private final ExponentialRetryAlgorithm recoveryAlgorithm; public ResumableUploadCallableImpl( ResumableUploadClient client, @@ -96,18 +97,16 @@ public ResumableUploadCallableImpl( this.retryingQueryCallable = createRetryingCallable( client.queryStatusCallable(), ResumableUploadCommand.QUERY, clientContext); - this.recoveryAlgorithm = - new ExponentialRetryAlgorithm(RETRY_SETTINGS, clientContext.getClock()); } @Override public ResumableUploadFuture futureCall( RequestT request, - InputStreamSupplier payload, + InputStreamSupplier payloadSupplier, @Nullable ApiCallContext context, @Nullable ResumableUploadOptions options) { checkNotNull(request, "request must not be null"); - checkNotNull(payload, "payload must not be null"); + checkNotNull(payloadSupplier, "payloadSupplier must not be null"); ResumableUploadOptions effectiveOptions = defaultOptions.merge(options); ApiCallContext effectiveCallContext = clientContext.getDefaultCallContext().merge(context); @@ -122,16 +121,25 @@ public ResumableUploadFuture futureCall( startFuture, retryingUploadChunkCallable, retryingQueryCallable, - payload, + payloadSupplier, effectiveOptions, - clientContext, - recoveryAlgorithm); + clientContext); } @Override public ResumableUploadFuture resumeCall( - String sessionUrl, InputStreamSupplier payload, @Nullable ResumableUploadOptions options) { - throw new UnsupportedOperationException("Session resumption is not yet implemented."); + String sessionUrl, + InputStreamSupplier payloadSupplier, + @Nullable ResumableUploadOptions options) { + checkArgument(!Strings.isNullOrEmpty(sessionUrl), "sessionUrl must not be null or empty"); + checkNotNull(payloadSupplier, "payloadSupplier must not be null"); + return ResumableUploadFutureImpl.createForResume( + sessionUrl, + retryingUploadChunkCallable, + retryingQueryCallable, + payloadSupplier, + defaultOptions.merge(options), + clientContext); } private static UnaryCallable createRetryingCallable( 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 c84a42d81a32..a5eb146fd2a5 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 @@ -47,6 +47,7 @@ import com.google.common.util.concurrent.MoreExecutors; import java.io.IOException; import java.io.InputStream; +import java.util.Optional; import java.util.concurrent.CancellationException; import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; @@ -85,9 +86,7 @@ final class ResumableUploadChunkCoordinator { String uploadUrl, InputStream payload, int chunkSize, - ApiCallContext callContext, - ExponentialRetryAlgorithm recoveryAlgorithm, - ScheduledExecutorService executor, + ClientContext clientContext, ResumableUploadProgressTracker progressTracker) { this.uploadChunkCallable = checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); @@ -96,9 +95,12 @@ final class ResumableUploadChunkCoordinator { this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null"); this.chunkSize = chunkSize; checkNotNull(payload, "payload must not be null"); - this.callContext = checkNotNull(callContext, "callContext must not be null"); - this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null"); - this.executor = checkNotNull(executor, "executor must not be null"); + checkNotNull(clientContext, "clientContext must not be null"); + this.callContext = clientContext.getDefaultCallContext(); + this.recoveryAlgorithm = + new ExponentialRetryAlgorithm( + ResumableUploadCallableImpl.RETRY_SETTINGS, clientContext.getClock()); + this.executor = checkNotNull(clientContext.getExecutor(), "executor must not be null"); this.recoverySettings = recoveryAlgorithm.createFirstAttempt(); this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null"); this.buffer = new RewindableStreamBuffer(payload, chunkSize); @@ -119,6 +121,15 @@ void start() { } } + void resume(QueryStatusResponse initialQueryResponse) { + registerCancellationListener(); + try { + handleQueryResponse(initialQueryResponse); + } catch (Throwable t) { + uploadResultFuture.setException(t); + } + } + private void registerCancellationListener() { uploadResultFuture.addListener( () -> { @@ -271,7 +282,10 @@ private void handleQueryResponse(QueryStatusResponse queryResponse) throw new IllegalStateException("Query status response missing X-Goog-Upload-Status header"); } if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) { - progressTracker.onFinalized(buffer.getBufferBaseOffset() + buffer.getPayload().length); + long finalizedBytes = + Optional.ofNullable(queryResponse.getCommittedOffset()) + .orElse(buffer.getBufferBaseOffset() + buffer.getPayload().length); + progressTracker.onFinalized(finalizedBytes); uploadResultFuture.set(queryResponse.getResponse()); return; } @@ -280,7 +294,7 @@ private void handleQueryResponse(QueryStatusResponse queryResponse) throw new IllegalStateException( "Incomplete query status response did not include a committed offset"); } - progressTracker.onOffsetReceived(committedOffset); + progressTracker.onOffsetReceived(committedOffset, chunkSize); buffer.realignTo(committedOffset); transmitChunk(); } 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 defc964090c3..c860a9f718ec 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 @@ -42,7 +42,6 @@ import com.google.api.gax.resumable.QueryStatusRequest; import com.google.api.gax.resumable.QueryStatusResponse; import com.google.api.gax.resumable.ResumableUploadSession; -import com.google.api.gax.retrying.ExponentialRetryAlgorithm; import com.google.common.util.concurrent.MoreExecutors; import com.google.errorprone.annotations.concurrent.GuardedBy; import java.io.IOException; @@ -55,6 +54,7 @@ import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.function.Consumer; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -71,15 +71,13 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur private final Object lock = new Object(); - private final ApiFuture startFuture; private final UnaryCallable> uploadChunkCallable; private final UnaryCallable> queryStatusCallable; private final ResumableUploadOptions options; - private final ApiCallContext callContext; + private final ClientContext clientContext; private final ScheduledExecutorService executor; - private final ExponentialRetryAlgorithm recoveryAlgorithm; private final ResumableUploadProgressTracker progressTracker; private final SettableApiFuture resultFuture = SettableApiFuture.create(); @@ -91,8 +89,9 @@ final class ResumableUploadFutureImpl implements ResumableUploadFutur @GuardedBy("lock") private @Nullable Integer effectiveChunkSize; - // Tracks the current operation's Future (start, chunk upload) to propagate cancellation, or - // null once a terminal transition (succeed, fail, cancel) has claimed completion. + // Tracks the current operation's Future (start, resume query, chunk upload) to propagate + // cancellation, or null once a terminal transition (succeed, fail, cancel) has claimed + // completion. @GuardedBy("lock") private @Nullable ApiFuture inFlightFuture; @@ -109,20 +108,73 @@ static ResumableUploadFutureImpl create( UnaryCallable> queryStatusCallable, InputStreamSupplier payloadSupplier, ResumableUploadOptions options, - ClientContext clientContext, - ExponentialRetryAlgorithm recoveryAlgorithm) { + ClientContext clientContext) { + checkNotNull(startFuture, "startFuture must not be null"); checkNotNull(payloadSupplier, "payloadSupplier must not be null"); ResumableUploadFutureImpl future = new ResumableUploadFutureImpl<>( - startFuture, + /* sessionUrl= */ null, uploadChunkCallable, queryStatusCallable, options, - clientContext, - recoveryAlgorithm, - ResumableUploadProgressTracker.create()); + clientContext); + try { + future.initiate( + startFuture, + payloadSupplier, + session -> { + ResumableUploadChunkCoordinator coordinator = + future.createAndRegisterCoordinator( + session.getUploadUrl(), session.getChunkGranularity()); + if (coordinator != null) { + coordinator.start(); + } + }); + } catch (Throwable t) { + future.fail(t); + } + return future; + } + + /** + * Creates and initiates a resumable upload future that resumes an existing upload session from + * {@code sessionUrl} by querying the server for its committed offset. + * + *

The stream opened from {@code payloadSupplier} is managed by the returned future and will be + * closed automatically upon completion, failure, or cancellation. + */ + static ResumableUploadFutureImpl createForResume( + String sessionUrl, + UnaryCallable> uploadChunkCallable, + UnaryCallable> queryStatusCallable, + InputStreamSupplier payloadSupplier, + ResumableUploadOptions options, + ClientContext clientContext) { + checkNotNull(sessionUrl, "sessionUrl must not be null"); + checkNotNull(payloadSupplier, "payloadSupplier must not be null"); + ResumableUploadFutureImpl future = + new ResumableUploadFutureImpl<>( + sessionUrl, uploadChunkCallable, queryStatusCallable, options, clientContext); + ApiFuture> queryFuture; try { - future.start(payloadSupplier); + queryFuture = + queryStatusCallable.futureCall( + QueryStatusRequest.create(sessionUrl), clientContext.getDefaultCallContext()); + } catch (Throwable t) { + queryFuture = ApiFutures.immediateFailedFuture(t); + } + try { + future.initiate( + queryFuture, + payloadSupplier, + queryResponse -> { + ResumableUploadChunkCoordinator coordinator = + future.createAndRegisterCoordinator( + sessionUrl, queryResponse.getChunkGranularity()); + if (coordinator != null) { + coordinator.resume(queryResponse); + } + }); } catch (Throwable t) { future.fail(t); } @@ -130,52 +182,51 @@ static ResumableUploadFutureImpl create( } private ResumableUploadFutureImpl( - ApiFuture startFuture, + @Nullable String sessionUrl, UnaryCallable> uploadChunkCallable, UnaryCallable> queryStatusCallable, ResumableUploadOptions options, - ClientContext clientContext, - ExponentialRetryAlgorithm recoveryAlgorithm, - ResumableUploadProgressTracker progressTracker) { - this.startFuture = checkNotNull(startFuture, "startFuture must not be null"); + ClientContext clientContext) { + this.uploadSessionUrl = sessionUrl; this.uploadChunkCallable = checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); this.queryStatusCallable = checkNotNull(queryStatusCallable, "queryStatusCallable must not be null"); this.options = checkNotNull(options, "options must not be null"); checkArgument(options.getChunkSize() > 0, "chunkSize must be > 0"); - checkNotNull(clientContext, "clientContext must not be null"); - this.callContext = clientContext.getDefaultCallContext(); + this.clientContext = checkNotNull(clientContext, "clientContext must not be null"); this.executor = checkNotNull(clientContext.getExecutor(), "executor must not be null"); - this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null"); - this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null"); - this.inFlightFuture = startFuture; + if (sessionUrl != null) { + this.progressTracker = ResumableUploadProgressTracker.createForResume(sessionUrl); + } else { + this.progressTracker = ResumableUploadProgressTracker.create(); + } ResumableUploadProgressListener progressListener = options.getProgressListener(); if (progressListener != null) { progressTracker.addListener(progressListener, executor); } } - private void start(InputStreamSupplier payloadSupplier) throws IOException { + private void initiate( + ApiFuture initiationFuture, + InputStreamSupplier payloadSupplier, + Consumer onSessionReady) + throws IOException { + // Claim inFlightFuture first so that fail() and cancel() can act on early failures. + synchronized (lock) { + this.inFlightFuture = initiationFuture; + } setUpGlobalTimeout(); this.payload = checkNotNull(payloadSupplier.get(), "payload supplier returned null"); ApiFutures.addCallback( - startFuture, - new ApiFutureCallback() { + initiationFuture, + new ApiFutureCallback() { @Override - public void onSuccess(ResumableUploadSession session) { - String sessionUrl = session.getUploadUrl(); - int chunkSize; + public void onSuccess(T result) { try { - chunkSize = alignChunkSize(options.getChunkSize(), session.getChunkGranularity()); + onSessionReady.accept(result); } catch (Throwable t) { fail(t); - return; - } - ResumableUploadChunkCoordinator coordinator = - createAndRegisterCoordinator(sessionUrl, chunkSize); - if (coordinator != null) { - coordinator.start(); } } @@ -198,7 +249,8 @@ private void setUpGlobalTimeout() { } private @Nullable ResumableUploadChunkCoordinator createAndRegisterCoordinator( - String sessionUrl, int chunkSize) { + String sessionUrl, int chunkGranularity) { + int chunkSize = alignChunkSize(options.getChunkSize(), chunkGranularity); ResumableUploadChunkCoordinator coordinator = new ResumableUploadChunkCoordinator<>( uploadChunkCallable, @@ -206,9 +258,7 @@ private void setUpGlobalTimeout() { sessionUrl, checkNotNull(payload), chunkSize, - callContext, - recoveryAlgorithm, - executor, + clientContext, progressTracker); ApiFuture uploadFuture = coordinator.getFuture(); synchronized (lock) { @@ -289,10 +339,12 @@ private ApiException toTerminalException(Throwable t) { return (ApiException) t; } String message = firstNonNull(t.getMessage(), t.getClass().getSimpleName()); - if (sessionUrl != null) { + if (sessionUrl != null && chunkSize != null) { message = String.format( "%s (upload URL: %s, effective chunk size: %d)", message, sessionUrl, chunkSize); + } else if (sessionUrl != null) { + message = String.format("%s (upload URL: %s)", message, sessionUrl); } // Rethrow with URL and effective chunk size included in message. if (t instanceof ApiException) { diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java index df2887effe20..054b0d96bcf2 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadProgressTracker.java @@ -101,6 +101,15 @@ static ResumableUploadProgressTracker create() { return new ResumableUploadProgressTracker(ResumableUploadProgress.newBuilder().build()); } + static ResumableUploadProgressTracker createForResume(String uploadUrl) { + checkNotNull(uploadUrl, "uploadUrl must not be null"); + return new ResumableUploadProgressTracker( + ResumableUploadProgress.newBuilder() + .setState(ResumableUploadProgress.STATE_RECOVERING) + .setUploadUrl(uploadUrl) + .build()); + } + void addListener(ResumableUploadProgressListener listener, Executor executor) { checkNotNull(listener, "listener must not be null"); checkNotNull(executor, "executor must not be null"); @@ -142,12 +151,13 @@ void onRecovering() { transition(builder -> builder.setState(ResumableUploadProgress.STATE_RECOVERING)); } - void onOffsetReceived(long committedOffset) { + void onOffsetReceived(long committedOffset, int effectiveChunkSize) { transition( builder -> builder .setState(ResumableUploadProgress.STATE_OFFSET_RECEIVED) - .setBytesUploaded(committedOffset)); + .setBytesUploaded(committedOffset) + .setEffectiveChunkSize(effectiveChunkSize)); } void onFinalized(long totalBytes) { diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java index 0030f0282793..06db44ca944b 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java @@ -33,6 +33,7 @@ import static com.google.common.base.Preconditions.checkNotNull; import com.google.common.io.ByteStreams; +import java.io.EOFException; import java.io.IOException; import java.io.InputStream; import java.util.Arrays; @@ -52,6 +53,7 @@ final class RewindableStreamBuffer { private int payloadLength; private boolean isFinal; private boolean streamExhausted; + private boolean hasReadFromStream; RewindableStreamBuffer(InputStream inputStream, int chunkSize) { this.inputStream = checkNotNull(inputStream, "inputStream must not be null"); @@ -70,6 +72,7 @@ final class RewindableStreamBuffer { * @throws IOException if reading from the stream fails */ void fill() throws IOException { + this.hasReadFromStream = true; this.bufferBaseOffset += payloadLength; this.payloadLength = ByteStreams.read(inputStream, buffer, 0, chunkSize); this.isFinal = (payloadLength < chunkSize); @@ -81,15 +84,39 @@ void fill() throws IOException { /** * Realigns the buffer window to {@code committedOffset}. * - *

Compacts forward within the existing buffer to discard already-committed bytes, and then - * tops up the buffer to capacity from the underlying stream. + *

Before the first {@link #fill()}, skips the underlying stream forward from byte {@code 0} to + * {@code committedOffset} and fills the first chunk. Once data has been read, compacts forward + * within the existing buffer to discard already-committed bytes and tops up the buffer to + * capacity from the underlying stream. * * @param committedOffset the server's committed byte offset - * @throws IllegalStateException if {@code committedOffset} is below the buffer's base offset or - * beyond the current buffer window - * @throws IOException if reading from the stream fails + * @throws IllegalStateException if {@code committedOffset} is negative, or (once data has been + * read) below the buffer's base offset or beyond the current buffer window + * @throws IOException if skipping or reading from the stream fails */ void realignTo(long committedOffset) throws IOException { + if (committedOffset < 0) { + throw new IllegalStateException( + String.format("Server committed offset %d must not be negative", committedOffset)); + } + if (!hasReadFromStream) { + try { + ByteStreams.skipFully(inputStream, committedOffset); + } catch (EOFException e) { + EOFException ex = + new EOFException( + String.format( + "Stream ended before reaching server committed offset %d; ensure the supplied" + + " stream is positioned at byte 0", + committedOffset)); + ex.initCause(e); + throw ex; + } + this.bufferBaseOffset = committedOffset; + this.payloadLength = 0; + fill(); + return; + } if (committedOffset < bufferBaseOffset) { throw new IllegalStateException( String.format( 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 6bcfc9d234a2..ad2eca4f1cec 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 @@ -465,13 +465,6 @@ void testUploadCallable_withContextAndSettings_appliesBoth() throws Exception { assertThat(future.get()).isEqualTo("done-both"); } - @Test - void testResumeCall_throwsUnsupportedOperationException() { - assertThrows( - UnsupportedOperationException.class, - () -> callable.resumeCall("https://upload.url/session", payloadOf("data"), null)); - } - @Test void testChunkRetry_transientFailureThenSuccess_retriesAndSucceeds() throws Exception { stubStartSession("https://upload.url/chunk-retry-ok"); @@ -1383,6 +1376,174 @@ public void close() throws IOException { .contains("stream close error"); } + @Test + void testResumeCall_partialOffset_resumesAndCompletes() throws Exception { + String sessionUrl = "https://upload.url/resume-partial"; + when(mockQueryCallable.futureCall(eq(QueryStatusRequest.create(sessionUrl)), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(8L, null, ResumableUploadStatus.ACTIVE))); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "resumed-done"))); + + RecordingProgressListener listener = new RecordingProgressListener(); + ResumableUploadOptions options = + defaultSettings.toBuilder().setProgressListener(listener).build(); + + TrackableStream stream = new TrackableStream("0123456789abcdefghij"); + ResumableUploadFuture future = callable.resumeCall(sessionUrl, () -> stream, options); + + assertThat(future.getUploadSessionUrl()).isEqualTo(sessionUrl); + assertThat(future.getEffectiveChunkSize()).isEqualTo(8); + assertThat(future.get()).isEqualTo("resumed-done"); + assertThat(stream.closed).isTrue(); + verifyNoInteractions(mockStartCallable); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any()); + List chunks = captor.getAllValues(); + assertChunk(chunks.get(0), 8, 8, false); + assertChunk(chunks.get(1), 16, 4, true); + assertThat(new String(chunks.get(1).getPayload(), StandardCharsets.UTF_8)).isEqualTo("ghij"); + + assertThat(listener.awaitTerminal()) + .containsExactly( + createProgress(ResumableUploadProgress.STATE_RECOVERING, 0L, sessionUrl, null), + createProgress(ResumableUploadProgress.STATE_OFFSET_RECEIVED, 8L, sessionUrl, 8), + createProgress(ResumableUploadProgress.STATE_UPLOADING, 16L, sessionUrl, 8), + createProgress(ResumableUploadProgress.STATE_FINALIZED, 20L, sessionUrl, 8)) + .inOrder(); + } + + @Test + void testResumeCall_alignsChunkSizeToQueryGranularity() throws Exception { + String sessionUrl = "https://upload.url/resume-granularity"; + when(mockQueryCallable.futureCall(eq(QueryStatusRequest.create(sessionUrl)), any())) + .thenReturn( + ApiFutures.immediateFuture( + QueryStatusResponse.newBuilder() + .setCommittedOffset(4L) + .setUploadStatus(ResumableUploadStatus.ACTIVE) + .setChunkGranularity(4) + .build())); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "resumed-aligned"))); + + // Configured chunkSize = 10, query granularity = 4 -> aligned down to 8 + ResumableUploadOptions options = defaultSettings.toBuilder().setChunkSize(10).build(); + ResumableUploadFuture future = + callable.resumeCall(sessionUrl, payloadOf("0123456789abcdef"), options); + + assertThat(future.get()).isEqualTo("resumed-aligned"); + assertThat(future.getEffectiveChunkSize()).isEqualTo(8); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any()); + List chunks = captor.getAllValues(); + assertChunk(chunks.get(0), 4, 8, false); + assertChunk(chunks.get(1), 12, 4, true); + } + + @Test + void testResumeCall_alreadyFinalizedOnServer_completesImmediately() throws Exception { + String sessionUrl = "https://upload.url/resume-already-final"; + when(mockQueryCallable.futureCall(eq(QueryStatusRequest.create(sessionUrl)), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(null, "already-finalized", ResumableUploadStatus.FINAL))); + + ResumableUploadFuture future = + callable.resumeCall(sessionUrl, payloadOf("0123456789abcdefghij"), null); + + assertThat(future.get()).isEqualTo("already-finalized"); + verifyNoInteractions(mockStartCallable); + verifyNoInteractions(mockChunkCallable); + } + + @Test + void testResumeCall_initialQueryTransientFailure_retriesAndResumes() throws Exception { + String sessionUrl = "https://upload.url/resume-query-retry"; + when(mockQueryCallable.futureCall(eq(QueryStatusRequest.create(sessionUrl)), any())) + .thenReturn( + ApiFutures.immediateFailedFuture(createApiException(503, StatusCode.Code.UNAVAILABLE))) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(8L, null, ResumableUploadStatus.ACTIVE))); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "resumed-after-503"))); + + ResumableUploadFuture future = + callable.resumeCall(sessionUrl, payloadOf("0123456789abcd"), null); + + assertThat(future.get()).isEqualTo("resumed-after-503"); + verify(mockQueryCallable, times(2)).futureCall(any(), any()); + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(1)).futureCall(captor.capture(), any()); + assertChunk(captor.getValue(), 8, 6, true); + assertThat(new String(captor.getValue().getPayload(), StandardCharsets.UTF_8)) + .isEqualTo("89abcd"); + } + + @Test + void testResumeCall_withSubsequentRecovery_realignsWithinChunkWindow() throws Exception { + String sessionUrl = "https://upload.url/resume-then-recover"; + // Initial resume query reports 8 bytes committed; mid-upload recovery reports 12 bytes + when(mockQueryCallable.futureCall(eq(QueryStatusRequest.create(sessionUrl)), any())) + .thenReturn( + ApiFutures.immediateFuture(createQueryResponse(8L, null, ResumableUploadStatus.ACTIVE))) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(12L, null, ResumableUploadStatus.ACTIVE))); + when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any())) + .thenReturn( + ApiFutures.immediateFailedFuture( + createApiException(400, StatusCode.Code.INVALID_ARGUMENT))) + .thenReturn( + ApiFutures.immediateFuture( + ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "recovered-after-resume"))); + + ResumableUploadFuture future = + callable.resumeCall(sessionUrl, payloadOf("0123456789abcdefgh"), null); + + assertThat(future.get()).isEqualTo("recovered-after-resume"); + ArgumentCaptor captor = ArgumentCaptor.forClass(ChunkUploadRequest.class); + verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any()); + List chunks = captor.getAllValues(); + assertChunk(chunks.get(0), 8, 8, false); + assertChunk(chunks.get(1), 12, 6, true); + assertThat(new String(chunks.get(1).getPayload(), StandardCharsets.UTF_8)).isEqualTo("cdefgh"); + } + + @Test + void testResumeCall_failure_includesSessionUrlAndEffectiveChunkSize() { + String sessionUrl = "https://upload.url/resume-short-stream"; + when(mockQueryCallable.futureCall(eq(QueryStatusRequest.create(sessionUrl)), any())) + .thenReturn( + ApiFutures.immediateFuture( + createQueryResponse(16L, null, ResumableUploadStatus.ACTIVE))); + + TrackableStream shortStream = new TrackableStream("01234567"); + ResumableUploadFuture future = callable.resumeCall(sessionUrl, () -> shortStream, null); + + ExecutionException ex = assertThrows(ExecutionException.class, future::get); + assertThat(ex.getCause()).isInstanceOf(UnknownException.class); + assertThat(ex.getCause().getMessage()).contains(sessionUrl); + assertThat(ex.getCause().getMessage()).contains("effective chunk size: 8"); + assertThat(shortStream.closed).isTrue(); + } + private static class HttpStatusStatusCode implements StatusCode { private final int httpStatus; private final StatusCode.Code code; diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java index 6ca4b8020b54..1aceab386cc8 100644 --- a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java @@ -33,6 +33,7 @@ import static org.junit.jupiter.api.Assertions.assertThrows; import java.io.ByteArrayInputStream; +import java.io.EOFException; import java.io.FilterInputStream; import java.io.IOException; import java.io.InputStream; @@ -205,4 +206,37 @@ public int read(byte[] b, int off, int len) throws IOException { assertThat(buffer.isFinal()).isFalse(); assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); } + + @Test + void testRealignTo_beforeFirstFill_skipsToCommittedOffset() throws IOException { + byte[] data = "0123456789ABCDEFGHIJ".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = new RewindableStreamBuffer(new ByteArrayInputStream(data), 8); + + buffer.realignTo(10L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(10L); + assertThat(buffer.getPayload()).hasLength(8); + assertThat(buffer.isFinal()).isFalse(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("ABCDEFGH"); + } + + @Test + void testRealignTo_beforeFirstFill_streamShorterThanOffset_throwsEofException() { + byte[] data = "01234".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = new RewindableStreamBuffer(new ByteArrayInputStream(data), 8); + + EOFException exception = assertThrows(EOFException.class, () -> buffer.realignTo(10L)); + assertThat(exception).hasMessageThat().contains("10"); + } + + @Test + void testRealignTo_beforeFirstFill_zeroOffset_fillsFromStart() throws IOException { + byte[] data = "0123456789".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = new RewindableStreamBuffer(new ByteArrayInputStream(data), 8); + + buffer.realignTo(0L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(0L); + assertThat(buffer.getPayload()).hasLength(8); + assertThat(buffer.isFinal()).isFalse(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); + } }