refactor(gax): invert resumable upload future and coordinator ownership

Author: whowes

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<ResponseT> {
   private final byte[] buffer;
   private final int chunkSize;
   private final ApiCallContext callContext;
-  private final ResumableUploadFutureImpl<ResponseT> sessionFuture;
+  private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
+  private volatile @Nullable ApiFuture<?> currentChunkFuture;
 
   ResumableUploadChunkCoordinator(
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
       String uploadUrl,
       InputStream payload,
       int chunkSize,
-      ApiCallContext callContext,
-      ResumableUploadFutureImpl<ResponseT> 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<ResponseT> 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,25 @@ private void transmitChunk(long currentOffset) {
     try {
       ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
           uploadChunkCallable.futureCall(chunkRequest, callContext);
-      sessionFuture.setInFlightFuture(chunkFuture);
+      this.currentChunkFuture = chunkFuture;
+      if (result.isCancelled()) {
+        chunkFuture.cancel(true);
+        return;
+      }
 
-      // 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()) {
+              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 +165,15 @@ public void onSuccess(ChunkUploadResponse<ResponseT> 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);
     }
   }
 }

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java

@@ -128,16 +128,46 @@ public void onSuccess(ResumableUploadSession session) {
             ResumableUploadChunkCoordinator<ResponseT> coordinator =
                 new ResumableUploadChunkCoordinator<>(
                     uploadChunkCallable,
-                    session.getUploadUrl(),
+                    uploadSessionUrl,
                     payload,
                     settings.getChunkSize(),
-                    callContext,
-                    ResumableUploadFutureImpl.this);
+                    callContext);
+            ApiFuture<ResponseT> uploadFuture;
             try {
-              coordinator.start();
+              uploadFuture = coordinator.start();
             } catch (Throwable t) {
               fail(t);
+              return;
+            }
+            boolean alreadyDone = false;
+            synchronized (lock) {
+              if (resultFuture.isDone()) {
+                alreadyDone = true;
+              } else {
+                inFlightFuture = uploadFuture;
+              }
+            }
+            if (alreadyDone) {
+              uploadFuture.cancel(true);
+              return;
             }
+            ApiFutures.addCallback(
+                uploadFuture,
+                new ApiFutureCallback<ResponseT>() {
+                  @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,33 +181,15 @@ 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;
     }
     closePayload();
     resultFuture.set(result);
   }
 
-  void fail(Throwable t) {
+  private void fail(Throwable t) {
     synchronized (lock) {
       inFlightFuture = null;
     }

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<ResumableUploadSession> startFuture = SettableApiFuture.create();
-    when(mockStartCallable.futureCall(any(), any())).thenReturn(startFuture);
-    ResumableUploadFuture<String> future =
-        callable.futureCall("resource-path", streamOf("data"), null);
-    assertThat(future.cancel(true)).isTrue();
-    assertThat(future.isCancelled()).isTrue();
-
-    SettableApiFuture<String> lateFuture = SettableApiFuture.create();
-    ((ResumableUploadFutureImpl<String>) future).setInFlightFuture(lateFuture);
-    assertThat(lateFuture.isCancelled()).isTrue();
-  }
-
   @Test
   void testUploadCallable_startFailure_failsFuture() {
     when(mockStartCallable.futureCall(any(), any()))