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

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 <ResponseT> the type of the final response message returned once the upload completes
  */
-@InternalApi
 @NullMarked
 final class ResumableUploadChunkCoordinator<ResponseT> {
 
   private static final byte[] EMPTY_PAYLOAD = new byte[0];
 
+  private final Object lock = new Object();
+
+  private final SettableApiFuture<ResponseT> result;
+  private final ApiFuture<ResumableUploadSession> startFuture;
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       uploadChunkCallable;
-  private final String uploadUrl;
   private final InputStream payload;
   private final byte[] buffer;
   private final int chunkSize;
   private final ApiCallContext callContext;
-  private final ResumableUploadFutureImpl<ResponseT> sessionFuture;
+
+  private volatile @Nullable String uploadSessionUrl;
+
+  @GuardedBy("lock")
+  private boolean done;
+
+  @GuardedBy("lock")
+  private boolean payloadClosed;
+
+  @GuardedBy("lock")
+  private @Nullable ApiFuture<?> inFlightFuture;
 
   ResumableUploadChunkCoordinator(
+      SettableApiFuture<ResponseT> result,
+      ApiFuture<ResumableUploadSession> startFuture,
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
-      String uploadUrl,
       InputStream payload,
-      int chunkSize,
-      ApiCallContext callContext,
-      ResumableUploadFutureImpl<ResponseT> 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<ResumableUploadSession>() {
+          @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<ChunkUploadResponse<ResponseT>> 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<ChunkUploadResponse<ResponseT>>() {
             @Override
             public void onSuccess(ChunkUploadResponse<ResponseT> 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<ResponseT> 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);
     }
   }
 }