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

@@ -75,6 +75,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
   private final long deadlineNanos;
   private final ApiCallContext callContext;
   private final ClientContext clientContext;
+  private final UploadProgressTracker progressTracker;
   private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
   private volatile @Nullable ApiFuture<?> currentChunkFuture;
 
@@ -85,7 +86,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
       InputStream payload,
       ResumableUploadCallSettings settings,
       long deadlineNanos,
-      ClientContext clientContext) {
+      ClientContext clientContext,
+      UploadProgressTracker progressTracker) {
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.queryStatusCallable =
@@ -96,6 +98,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
     this.deadlineNanos = deadlineNanos;
     this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
     this.callContext = clientContext.getDefaultCallContext();
+    this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
     this.buffer = new RewindableStreamBuffer(payload, settings.getChunkSize(), uploadUrl);
   }
 
@@ -145,7 +148,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
                 return;
               }
               if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
-                recover();
+                recover(null);
               } else {
                 onChunkUploaded(response);
               }
@@ -159,7 +162,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.
@@ -173,7 +176,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(cause);
     try {
       // Dispatch the query status call and register the in-flight future for cancellation.
       ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
@@ -237,22 +247,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,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.BetaApi;
+import java.util.concurrent.Executor;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -48,4 +49,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

@@ -76,11 +76,15 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   private final ResumableUploadCallSettings settings;
   private final ClientContext clientContext;
   private final ScheduledExecutorService executor;
+  private final UploadProgressTracker progressTracker = new UploadProgressTracker();
   private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
 
   private volatile @Nullable String uploadSessionUrl;
   private volatile long deadlineNanos;
 
+  @GuardedBy("lock")
+  private boolean done;
+
   @GuardedBy("lock")
   private @Nullable ApiFuture<?> inFlightFuture;
 
@@ -149,10 +153,16 @@ 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,
@@ -161,7 +171,8 @@ public void onSuccess(ResumableUploadSession session) {
                     payload,
                     settings,
                     deadlineNanos,
-                    clientContext);
+                    clientContext,
+                    progressTracker);
             ApiFuture<ResponseT> uploadFuture;
             try {
               uploadFuture = coordinator.start();
@@ -171,7 +182,7 @@ public void onSuccess(ResumableUploadSession session) {
             }
             boolean alreadyDone = false;
             synchronized (lock) {
-              if (resultFuture.isDone()) {
+              if (done) {
                 alreadyDone = true;
               } else {
                 inFlightFuture = uploadFuture;
@@ -224,13 +235,18 @@ private void onTimeout() {
   private void succeed(@Nullable ResponseT result) {
     ScheduledFuture<?> timeout;
     synchronized (lock) {
+      if (done) {
+        return;
+      }
+      done = true;
       inFlightFuture = null;
       timeout = this.timeoutFuture;
       this.timeoutFuture = null;
     }
     if (timeout != null) {
       timeout.cancel(false);
     }
+    progressTracker.onFinalized(progressTracker.getStatus().getBytesUploaded());
     closePayload();
     resultFuture.set(result);
   }
@@ -239,6 +255,10 @@ private void fail(Throwable t) {
     ScheduledFuture<?> timeout;
     ApiFuture<?> inFlight;
     synchronized (lock) {
+      if (done) {
+        return;
+      }
+      done = true;
       inFlight = this.inFlightFuture;
       this.inFlightFuture = null;
       timeout = this.timeoutFuture;
@@ -250,6 +270,7 @@ private void fail(Throwable t) {
     if (inFlight != null) {
       inFlight.cancel(true);
     }
+    progressTracker.onFailed(t, uploadSessionUrl);
     closePayload();
     resultFuture.setException(t);
   }
@@ -267,6 +288,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);
@@ -278,6 +309,10 @@ public boolean cancel(boolean mayInterruptIfRunning) {
     ApiFuture<?> inFlight;
     ScheduledFuture<?> timeout;
     synchronized (lock) {
+      if (done) {
+        return false;
+      }
+      done = true;
       cancelled = resultFuture.cancel(mayInterruptIfRunning);
       inFlight = this.inFlightFuture;
       this.inFlightFuture = null;
@@ -290,6 +325,7 @@ public boolean cancel(boolean mayInterruptIfRunning) {
     if (inFlight != null) {
       inFlight.cancel(mayInterruptIfRunning);
     }
+    progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
     closePayload();
     return cancelled;
   }