feat(gax): wire progress listener into coordinator and attempt lifecycle

Author: whowes

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

@@ -70,6 +70,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
   private final String uploadUrl;
   private final RewindableStreamBuffer buffer;
   private final ApiCallContext callContext;
+  private final ResumableUploadProgressTracker progressTracker;
   private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
   private volatile @Nullable ApiFuture<?> currentChunkFuture;
 
@@ -79,14 +80,16 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
       String uploadUrl,
       InputStream payload,
       int chunkSize,
-      ApiCallContext callContext) {
+      ApiCallContext callContext,
+      ResumableUploadProgressTracker progressTracker) {
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.queryStatusCallable =
         checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
     this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
     checkNotNull(payload, "payload must not be null");
     this.callContext = checkNotNull(callContext, "callContext must not be null");
+    this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
     this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl);
   }
 
@@ -136,7 +139,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
                 return;
               }
               if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
-                recover();
+                recover(null);
               } else {
                 onChunkUploaded(response);
               }
@@ -150,7 +153,7 @@ public void onFailure(Throwable t) {
               Category category =
                   ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD);
               if (category == Category.RECOVERABLE) {
-                recover();
+                recover(t);
               } else {
                 // Category.TRANSIENT errors reaching here have already exhausted their retry budget
                 // in the underlying RetryingCallable and become fatal per protocol specification.
@@ -164,7 +167,14 @@ public void onFailure(Throwable t) {
     }
   }
 
-  private void recover() {
+  /**
+   * Queries the session for the server's committed offset and resumes transmission from it.
+   *
+   * @param cause the error that triggered recovery, or null if the server acknowledged the chunk
+   *     without an upload status header
+   */
+  private void recover(@Nullable Throwable cause) {
+    progressTracker.onRecovering();
     try {
       // Dispatch the query status call and register the in-flight future for cancellation.
       ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
@@ -228,22 +238,29 @@ private void handleQueryResponse(QueryStatusResponse<ResponseT> queryResponse)
           "Incomplete query status response did not include a committed offset for upload URL: "
               + uploadUrl);
     }
+    progressTracker.onOffsetReceived(committedOffset);
     buffer.realignTo(committedOffset);
     dispatchCurrentChunk();
   }
 
   private void onChunkUploaded(ChunkUploadResponse<ResponseT> response) {
-    long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
-    if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
-      result.set(response.getResponse());
-    } else if (buffer.isFinal()) {
-      result.setException(
-          new IllegalStateException(
-              "Upload stream ended and final chunk was transmitted, but server returned"
-                  + " incomplete status for upload URL: "
-                  + uploadUrl));
-    } else {
-      chunkExecutor.execute(() -> transmitChunk(nextOffset));
+    try {
+      long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
+      if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
+        progressTracker.onChunkUploaded(nextOffset);
+        result.set(response.getResponse());
+      } else if (buffer.isFinal()) {
+        result.setException(
+            new IllegalStateException(
+                "Upload stream ended and final chunk was transmitted, but server returned"
+                    + " incomplete status for upload URL: "
+                    + uploadUrl));
+      } else {
+        progressTracker.onChunkUploaded(nextOffset);
+        chunkExecutor.execute(() -> transmitChunk(nextOffset));
+      }
+    } catch (Throwable t) {
+      result.setException(t);
     }
   }
 

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

@@ -31,6 +31,9 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.BetaApi;
+import com.google.api.gax.resumable.ResumableUploadProgress;
+import com.google.api.gax.resumable.ResumableUploadProgressListener;
+import java.util.concurrent.Executor;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -48,4 +51,18 @@ public interface ResumableUploadFuture<ResponseT> extends ApiFuture<ResponseT> {
 
   /** Returns the upload session URL, or {@code null} if session initiation is in progress. */
   @Nullable String getUploadSessionUrl();
+
+  /**
+   * Registers a listener to receive progress and state transition notifications for this upload.
+   *
+   * <p>A snapshot of the current upload status is dispatched to the listener immediately upon
+   * subscription on the provided executor. Subsequent status updates are delivered in order.
+   *
+   * @param listener callback listener to receive progress notifications
+   * @param executor executor on which the listener callbacks are dispatched
+   */
+  void addProgressListener(ResumableUploadProgressListener listener, Executor executor);
+
+  /** Returns the current progress snapshot of the upload session. */
+  ResumableUploadProgress getStatus();
 }

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

@@ -41,6 +41,8 @@
 import com.google.api.gax.resumable.ChunkUploadResponse;
 import com.google.api.gax.resumable.QueryStatusRequest;
 import com.google.api.gax.resumable.QueryStatusResponse;
+import com.google.api.gax.resumable.ResumableUploadProgress;
+import com.google.api.gax.resumable.ResumableUploadProgressListener;
 import com.google.api.gax.resumable.ResumableUploadSession;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
@@ -58,10 +60,11 @@
 import org.jspecify.annotations.Nullable;
 
 /**
- * Implementation of {@link ResumableUploadFuture} responsible for the end-to-end management of a
- * resumable upload session.
+ * Baseline implementation of {@link ResumableUploadFuture}.
  *
- * @param <ResponseT> the type of the final response message returned once the upload completes
+ * <p>Coordinates session initiation, delegates chunk streaming to {@link
+ * ResumableUploadChunkCoordinator}, and manages the lifecycle of the user-provided payload {@link
+ * InputStream}.
  */
 @NullMarked
 final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFuture<ResponseT> {
@@ -79,10 +82,14 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   private final ResumableUploadCallSettings settings;
   private final ApiCallContext callContext;
   private final ScheduledExecutorService executor;
+  private final ResumableUploadProgressTracker progressTracker = new ResumableUploadProgressTracker();
   private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
 
   private volatile @Nullable String uploadSessionUrl;
 
+  @GuardedBy("lock")
+  private boolean done;
+
   @GuardedBy("lock")
   private @Nullable ApiFuture<?> inFlightFuture;
 
@@ -147,18 +154,25 @@ private void start() {
         new ApiFutureCallback<ResumableUploadSession>() {
           @Override
           public void onSuccess(ResumableUploadSession session) {
+            synchronized (lock) {
+              if (done) {
+                return;
+              }
+              uploadSessionUrl = session.getUploadUrl();
+            }
+            progressTracker.onStarted(uploadSessionUrl);
             if (resultFuture.isDone()) {
               return;
             }
-            uploadSessionUrl = session.getUploadUrl();
             ResumableUploadChunkCoordinator<ResponseT> coordinator =
                 new ResumableUploadChunkCoordinator<>(
                     uploadChunkCallable,
                     queryStatusCallable,
                     uploadSessionUrl,
                     payload,
                     settings.getChunkSize(),
-                    callContext);
+                    callContext,
+                    progressTracker);
             ApiFuture<ResponseT> uploadFuture;
             try {
               uploadFuture = coordinator.start();
@@ -168,7 +182,7 @@ public void onSuccess(ResumableUploadSession session) {
             }
             boolean alreadyDone = false;
             synchronized (lock) {
-              if (resultFuture.isDone()) {
+              if (done) {
                 alreadyDone = true;
               } else {
                 inFlightFuture = uploadFuture;
@@ -221,21 +235,31 @@ private void onTimeout() {
 
   private void succeed(@Nullable ResponseT result) {
     synchronized (lock) {
+      if (done) {
+        return;
+      }
+      done = true;
       inFlightFuture = null;
     }
+    progressTracker.onFinalized(progressTracker.getStatus().getBytesUploaded());
     closePayload();
     resultFuture.set(result);
   }
 
   private void fail(Throwable t) {
     ApiFuture<?> inFlight;
     synchronized (lock) {
+      if (done) {
+        return;
+      }
+      done = true;
       inFlight = this.inFlightFuture;
       this.inFlightFuture = null;
     }
     if (inFlight != null) {
       inFlight.cancel(true);
     }
+    progressTracker.onFailed();
     closePayload();
     resultFuture.setException(t);
   }
@@ -253,6 +277,16 @@ private void closePayload() {
     return uploadSessionUrl;
   }
 
+  @Override
+  public void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
+    progressTracker.addListener(listener, executor);
+  }
+
+  @Override
+  public ResumableUploadProgress getStatus() {
+    return progressTracker.getStatus();
+  }
+
   @Override
   public void addListener(Runnable listener, Executor executor) {
     resultFuture.addListener(listener, executor);
@@ -263,13 +297,18 @@ public boolean cancel(boolean mayInterruptIfRunning) {
     boolean cancelled;
     ApiFuture<?> inFlight;
     synchronized (lock) {
+      if (done) {
+        return false;
+      }
+      done = true;
       cancelled = resultFuture.cancel(mayInterruptIfRunning);
       inFlight = this.inFlightFuture;
       this.inFlightFuture = null;
     }
     if (inFlight != null) {
       inFlight.cancel(mayInterruptIfRunning);
     }
+    progressTracker.onFailed();
     closePayload();
     return cancelled;
   }