From 572db5bceafd30d7fae7af9d370ffbeb61954312 Mon Sep 17 00:00:00 2001 From: whowes Date: Mon, 14 Sep 2026 17:52:25 +0000 Subject: [PATCH] refactor(gax): invert resumable upload future and coordinator ownership --- .../rpc/ResumableUploadChunkCoordinator.java | 177 +++++++++++++++--- .../gax/rpc/ResumableUploadFutureImpl.java | 172 +++-------------- .../rpc/ResumableUploadCallableImplTest.java | 47 ----- 3 files changed, 170 insertions(+), 226 deletions(-) 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 2d59f89342f4..0703b16a9b7b 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 @@ -29,79 +29,189 @@ */ 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.InternalApi; +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.common.io.ByteStreams; import com.google.common.util.concurrent.MoreExecutors; +import com.google.errorprone.annotations.concurrent.GuardedBy; import java.io.IOException; import java.io.InputStream; import java.util.Arrays; import java.util.concurrent.CancellationException; import org.jspecify.annotations.NullMarked; +import org.jspecify.annotations.Nullable; /** - * Coordinates chunk transmission steps of a resumable upload session. + * Coordinates chunk transmission steps and manages lifecycle 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 Object lock = new Object(); + + private final SettableApiFuture result; + private final ApiFuture startFuture; private final UnaryCallable> uploadChunkCallable; - private final String uploadUrl; private final InputStream payload; private final byte[] buffer; private final int chunkSize; private final ApiCallContext callContext; - private final ResumableUploadFutureImpl sessionFuture; + + private volatile @Nullable String uploadSessionUrl; + + @GuardedBy("lock") + private boolean done; + + @GuardedBy("lock") + private boolean payloadClosed; + + @GuardedBy("lock") + private @Nullable ApiFuture inFlightFuture; ResumableUploadChunkCoordinator( + SettableApiFuture result, + ApiFuture startFuture, UnaryCallable> uploadChunkCallable, - String uploadUrl, InputStream payload, - int chunkSize, - ApiCallContext callContext, - ResumableUploadFutureImpl sessionFuture) { + ResumableUploadCallSettings settings, + ApiCallContext callContext) { + this.result = checkNotNull(result, "result must not be null"); + this.startFuture = checkNotNull(startFuture, "startFuture must not be null"); this.uploadChunkCallable = checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null"); - this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null"); this.payload = checkNotNull(payload, "payload must not be null"); - this.chunkSize = chunkSize; + checkNotNull(settings, "settings must not be null"); + checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0"); + this.chunkSize = settings.getChunkSize(); this.callContext = checkNotNull(callContext, "callContext must not be null"); - this.sessionFuture = checkNotNull(sessionFuture, "sessionFuture must not be null"); this.buffer = new byte[chunkSize]; + synchronized (lock) { + this.inFlightFuture = startFuture; + } } void start() { - transmitChunk(0L); + ApiFutures.addCallback( + startFuture, + new ApiFutureCallback() { + @Override + public void onSuccess(ResumableUploadSession session) { + synchronized (lock) { + if (done) { + return; + } + } + uploadSessionUrl = session.getUploadUrl(); + transmitChunk(0L); + } + + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException) { + return; + } + finish(null, t); + } + }, + MoreExecutors.directExecutor()); + } + + @Nullable String getUploadSessionUrl() { + return uploadSessionUrl; + } + + void setInFlightFuture(ApiFuture future) { + boolean shouldCancel = false; + synchronized (lock) { + if (done) { + shouldCancel = result.isCancelled(); + } else { + this.inFlightFuture = future; + } + } + if (shouldCancel) { + future.cancel(true); + } + } + + void cancel(boolean mayInterruptIfRunning) { + ApiFuture inFlight; + synchronized (lock) { + if (done) { + return; + } + done = true; + inFlight = this.inFlightFuture; + this.inFlightFuture = null; + } + if (inFlight != null) { + inFlight.cancel(mayInterruptIfRunning); + } + closePayload(); + } + + private void finish(@Nullable ResponseT response, @Nullable Throwable error) { + synchronized (lock) { + if (done) { + return; + } + done = true; + inFlightFuture = null; + } + IOException closeError = closePayload(); + if (error == null) { + result.set(response); + } else { + if (closeError != null) { + error.addSuppressed(closeError); + } + result.setException(error); + } + } + + private @Nullable IOException closePayload() { + synchronized (lock) { + if (payloadClosed) { + return null; + } + payloadClosed = true; + } + try { + payload.close(); + return null; + } catch (IOException e) { + return e; + } } private void transmitChunk(long currentOffset) { - // Abort if the session was already completed or canceled. - if (sessionFuture.isDone()) { - return; + synchronized (lock) { + if (done) { + return; + } } - // Read the next chunk slice from the payload stream. int bytesRead; try { bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize); } catch (IOException e) { - sessionFuture.fail(e); + finish(null, e); return; } - // Determine if this is the final chunk and build the chunk request. boolean isFinal = bytesRead < chunkSize; byte[] chunkPayload; if (bytesRead == chunkSize) { @@ -112,35 +222,42 @@ private void transmitChunk(long currentOffset) { chunkPayload = Arrays.copyOf(buffer, bytesRead); } + String url = uploadSessionUrl; + if (url == null) { + finish(null, new IllegalStateException("Upload session URL not available")); + return; + } + ChunkUploadRequest chunkRequest = ChunkUploadRequest.newBuilder() - .setUploadUrl(uploadUrl) + .setUploadUrl(url) .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); + 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; + synchronized (lock) { + if (done) { + return; + } } long nextOffset = currentOffset + chunkLength; if (response.isComplete()) { - sessionFuture.succeed(response.getResponse()); + finish(response.getResponse(), null); } else if (isFinal) { - sessionFuture.fail( + finish( + null, new IllegalStateException( "Upload stream ended and final chunk was transmitted, but server returned" + " incomplete status")); @@ -151,15 +268,15 @@ public void onSuccess(ChunkUploadResponse response) { @Override public void onFailure(Throwable t) { - if (t instanceof CancellationException || sessionFuture.isDone()) { + if (t instanceof CancellationException) { return; } - sessionFuture.fail(t); + finish(null, t); } }, MoreExecutors.directExecutor()); } catch (Throwable t) { - sessionFuture.fail(t); + finish(null, t); } } } 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 b37f14d853c0..094727707e07 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java @@ -29,21 +29,14 @@ */ 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.common.util.concurrent.MoreExecutors; -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.TimeUnit; @@ -52,191 +45,72 @@ import org.jspecify.annotations.Nullable; /** - * Implementation of {@link ResumableUploadFuture} responsible for the end-to-end management of a - * resumable upload session. + * Caller-facing delegation handle for 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 SettableApiFuture result; + private final ResumableUploadChunkCoordinator coordinator; - private final ApiFuture startFuture; - private final UnaryCallable> - uploadChunkCallable; - private final InputStream payload; - private final ResumableUploadCallSettings settings; - private final ApiCallContext callContext; - 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) { - ResumableUploadFutureImpl future = - new ResumableUploadFutureImpl<>( - startFuture, uploadChunkCallable, payload, settings, callContext); - try { - future.start(); - } catch (Throwable t) { - future.fail(t); - } - return future; - } - - private ResumableUploadFutureImpl( - ApiFuture startFuture, - UnaryCallable> uploadChunkCallable, - InputStream payload, - ResumableUploadCallSettings settings, - ApiCallContext callContext) { - 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.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, - 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); - } - }, - MoreExecutors.directExecutor()); - } - - /** - * 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); + SettableApiFuture result = SettableApiFuture.create(); + ResumableUploadChunkCoordinator coordinator = + new ResumableUploadChunkCoordinator<>( + result, startFuture, uploadChunkCallable, payload, settings, callContext); + ResumableUploadFutureImpl handle = + new ResumableUploadFutureImpl<>(result, coordinator); + coordinator.start(); + return handle; } - private void closePayload() { - try { - payload.close(); - } catch (IOException ignored) { - // Suppressed during stream cleanup - } + ResumableUploadFutureImpl( + SettableApiFuture result, ResumableUploadChunkCoordinator coordinator) { + this.result = checkNotNull(result, "result must not be null"); + this.coordinator = checkNotNull(coordinator, "coordinator must not be null"); } @Override public @Nullable String getUploadSessionUrl() { - return uploadSessionUrl; + return coordinator.getUploadSessionUrl(); } @Override public void addListener(Runnable listener, Executor executor) { - resultFuture.addListener(listener, executor); + result.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; + coordinator.cancel(mayInterruptIfRunning); + return result.cancel(mayInterruptIfRunning); } @Override public boolean isCancelled() { - return resultFuture.isCancelled(); + return result.isCancelled(); } @Override public boolean isDone() { - return resultFuture.isDone(); + return result.isDone(); } @Override public ResponseT get() throws InterruptedException, ExecutionException { - return resultFuture.get(); + return result.get(); } @Override public ResponseT get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { - return resultFuture.get(timeout, unit); + return result.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 index 1abbedd55a77..3654796f83fe 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 @@ -214,20 +214,6 @@ void testUploadCallable_cancelInFlight_haltsUpload() throws Exception { 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())) @@ -280,39 +266,6 @@ void testUploadCallable_closesPayloadOnFailure() { 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_withApiCallContext_mergesAndPassesContext() throws Exception { stubStartSession("https://upload.url/context");