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 5ea5723d483c..89044fd8ff42 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 @@ -35,6 +35,7 @@ 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.ResumableUploadStatus; @@ -45,6 +46,7 @@ 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. @@ -64,32 +66,40 @@ final class ResumableUploadChunkCoordinator { private final byte[] buffer; private final int chunkSize; private final ApiCallContext callContext; - private final ResumableUploadFutureImpl sessionFuture; + private final SettableApiFuture result = SettableApiFuture.create(); + private volatile @Nullable ApiFuture currentChunkFuture; ResumableUploadChunkCoordinator( UnaryCallable> uploadChunkCallable, String uploadUrl, InputStream payload, int chunkSize, - ApiCallContext callContext, - ResumableUploadFutureImpl sessionFuture) { + ApiCallContext callContext) { 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.sessionFuture = checkNotNull(sessionFuture, "sessionFuture must not be null"); this.buffer = new byte[chunkSize]; } - void start() { + ApiFuture start() { + result.addListener( + () -> { + ApiFuture chunk = currentChunkFuture; + if (result.isCancelled() && chunk != null) { + chunk.cancel(true); + } + }, + MoreExecutors.directExecutor()); transmitChunk(0L); + return result; } private void transmitChunk(long currentOffset) { // Abort if the session was already completed or canceled. - if (sessionFuture.isDone()) { + if (result.isDone()) { return; } @@ -98,7 +108,7 @@ private void transmitChunk(long currentOffset) { try { bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize); } catch (IOException e) { - sessionFuture.fail(e); + result.setException(e); return; } @@ -126,22 +136,21 @@ private void transmitChunk(long currentOffset) { try { ApiFuture> chunkFuture = uploadChunkCallable.futureCall(chunkRequest, callContext); - sessionFuture.setInFlightFuture(chunkFuture); + this.currentChunkFuture = 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()) { + if (result.isDone()) { return; } long nextOffset = currentOffset + chunkLength; if (response.getUploadStatus() == ResumableUploadStatus.FINAL) { - sessionFuture.succeed(response.getResponse()); + result.set(response.getResponse()); } else if (isFinal) { - sessionFuture.fail( + result.setException( new IllegalStateException( "Upload stream ended and final chunk was transmitted, but server returned" + " incomplete status")); @@ -152,15 +161,15 @@ public void onSuccess(ChunkUploadResponse response) { @Override public void onFailure(Throwable t) { - if (t instanceof CancellationException || sessionFuture.isDone()) { + if (t instanceof CancellationException || result.isDone()) { return; } - sessionFuture.fail(t); + result.setException(t); } }, MoreExecutors.directExecutor()); } catch (Throwable t) { - sessionFuture.fail(t); + result.setException(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..644b52e30777 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 @@ -128,16 +128,40 @@ public void onSuccess(ResumableUploadSession session) { ResumableUploadChunkCoordinator coordinator = new ResumableUploadChunkCoordinator<>( uploadChunkCallable, - session.getUploadUrl(), + uploadSessionUrl, payload, settings.getChunkSize(), - callContext, - ResumableUploadFutureImpl.this); + callContext); + ApiFuture uploadFuture; try { - coordinator.start(); + uploadFuture = coordinator.start(); } catch (Throwable t) { fail(t); + return; + } + synchronized (lock) { + if (resultFuture.isDone()) { + return; + } + inFlightFuture = uploadFuture; } + ApiFutures.addCallback( + uploadFuture, + new ApiFutureCallback() { + @Override + public void onSuccess(ResponseT response) { + succeed(response); + } + + @Override + public void onFailure(Throwable t) { + if (t instanceof CancellationException) { + return; + } + fail(t); + } + }, + MoreExecutors.directExecutor()); } @Override @@ -151,25 +175,7 @@ public void onFailure(Throwable 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) { + private void succeed(@Nullable ResponseT result) { synchronized (lock) { inFlightFuture = null; } @@ -177,7 +183,7 @@ void succeed(@Nullable ResponseT result) { resultFuture.set(result); } - void fail(Throwable t) { + private void fail(Throwable t) { synchronized (lock) { inFlightFuture = null; } diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java index 66932ebaf8f6..b882a43df401 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 @@ -229,20 +229,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()))