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/ChunkAttemptCallable.java

@@ -72,6 +72,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
   private final RewindableStreamBuffer buffer;
   private final String uploadUrl;
   private final ApiCallContext originalCallContext;
+  private final UploadProgressTracker progressTracker;
   private final long deadlineNanos;
   private final ApiClock clock;
 
@@ -90,7 +91,8 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
       String uploadUrl,
       ChunkUploadRequest request,
       ApiCallContext callContext,
-      ResumableUploadCommand command) {
+      ResumableUploadCommand command,
+      UploadProgressTracker progressTracker) {
     this(
         uploadChunkCallable,
         queryStatusCallable,
@@ -99,6 +101,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
         request,
         callContext,
         command,
+        progressTracker,
         Long.MAX_VALUE,
         NanoClock.getDefaultClock());
   }
@@ -111,6 +114,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
       ChunkUploadRequest request,
       ApiCallContext callContext,
       ResumableUploadCommand command,
+      UploadProgressTracker progressTracker,
       long deadlineNanos,
       ApiClock clock) {
     this.uploadChunkCallable =
@@ -122,6 +126,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
     this.currentRequest = checkNotNull(request, "request must not be null");
     this.originalCallContext = checkNotNull(callContext, "callContext must not be null");
     this.currentCommand = checkNotNull(command, "command must not be null");
+    this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
     this.deadlineNanos = deadlineNanos;
     this.clock = checkNotNull(clock, "clock must not be null");
   }
@@ -164,6 +169,8 @@ private void prepareAttempt(
       ApiCallContext attemptContext,
       RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
     this.currentCommand = ResumableUploadCommand.QUERY;
+    progressTracker.onRecovering(lastFailure);
+
     // Per GAX-R7: query uses sensible unary defaults trimmed to the remaining global deadline.
     long remainingNanos =
         deadlineNanos == Long.MAX_VALUE
@@ -259,6 +266,7 @@ private void handleQuerySuccess(
     // Normal path: realign buffer to committedOffset, compact and top up.
     try {
       buffer.realignTo(committedOffset);
+      progressTracker.onOffsetReceived(committedOffset);
     } catch (Throwable e) {
       failAttempt(attemptFuture, e);
       return;

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

@@ -85,6 +85,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;
 
@@ -96,7 +97,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
       ResumableUploadCallSettings settings,
       long deadlineNanos,
       ApiCallContext callContext,
-      ClientContext clientContext) {
+      ClientContext clientContext,
+      UploadProgressTracker progressTracker) {
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.queryStatusCallable =
@@ -107,6 +109,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
     this.deadlineNanos = deadlineNanos;
     this.callContext = checkNotNull(callContext, "callContext must not be null");
     this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
+    this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
     this.buffer = new RewindableStreamBuffer(payload, settings.getChunkSize(), uploadUrl);
   }
 
@@ -198,6 +201,7 @@ private void transmitSingleChunk(long currentOffset) {
             chunkRequest,
             chunkCallContext,
             command,
+            progressTracker,
             deadlineNanos,
             clientContext.getClock());
 
@@ -222,6 +226,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
               return;
             }
             long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
+            progressTracker.onChunkUploaded(nextOffset);
             if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
               result.set(response.getResponse());
             } else if (buffer.isFinal()) {

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

@@ -77,6 +77,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   private final ApiCallContext callContext;
   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;
@@ -160,6 +161,7 @@ public void onSuccess(ResumableUploadSession session) {
                 return;
               }
               uploadSessionUrl = session.getUploadUrl();
+              progressTracker.onStarted(uploadSessionUrl);
               ResumableUploadChunkCoordinator<ResponseT> coordinator =
                   new ResumableUploadChunkCoordinator<>(
                       uploadChunkCallable,
@@ -169,7 +171,8 @@ public void onSuccess(ResumableUploadSession session) {
                       settings,
                       deadlineNanos,
                       callContext,
-                      clientContext);
+                      clientContext,
+                      progressTracker);
               uploadFuture = coordinator.start();
               inFlightFuture = uploadFuture;
             }
@@ -241,11 +244,13 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) {
     }
     IOException closeError = closePayload();
     if (error == null) {
+      progressTracker.onFinalized(progressTracker.getStatus().getBytesUploaded());
       resultFuture.set(response);
     } else {
       if (closeError != null) {
         error.addSuppressed(closeError);
       }
+      progressTracker.onFailed(error, uploadSessionUrl);
       resultFuture.setException(error);
     }
   }
@@ -270,6 +275,16 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) {
     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);
@@ -295,6 +310,7 @@ public boolean cancel(boolean mayInterruptIfRunning) {
     if (inFlight != null) {
       inFlight.cancel(mayInterruptIfRunning);
     }
+    progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
     closePayload();
     return resultFuture.cancel(mayInterruptIfRunning);
   }

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ChunkAttemptCallableTest.java

@@ -131,7 +131,8 @@ void call_successfulChunk_setsAttemptFuture() throws Exception {
             "https://upload.url/test",
             request,
             callContext,
-            ResumableUploadCommand.UPLOAD);
+            ResumableUploadCommand.UPLOAD,
+            new UploadProgressTracker());
 
     callable.setRetryingFuture(mockExternalFuture);
     ChunkUploadResponse<String> callResult = callable.call();
@@ -179,7 +180,8 @@ void call_returnsWithoutBlocking_andPropagatesCancellation() throws Exception {
             "https://upload.url/test",
             request,
             callContext,
-            ResumableUploadCommand.UPLOAD);
+            ResumableUploadCommand.UPLOAD,
+            new UploadProgressTracker());
 
     List<Runnable> listeners = new ArrayList<>();
     doAnswer(
@@ -253,7 +255,8 @@ void call_perAttemptDeadline_appliesRpcTimeoutToCallContext() throws Exception {
             "https://upload.url/test",
             request,
             callContext,
-            ResumableUploadCommand.UPLOAD);
+            ResumableUploadCommand.UPLOAD,
+            new UploadProgressTracker());
 
     callable.setRetryingFuture(mockExternalFuture);
     callable.call();
@@ -294,7 +297,8 @@ void call_nonBlockingExecution_callingThreadMakesImmediateProgress() throws Exce
             "https://upload.url/test",
             request,
             callContext,
-            ResumableUploadCommand.UPLOAD);
+            ResumableUploadCommand.UPLOAD,
+            new UploadProgressTracker());
 
     callable.setRetryingFuture(mockExternalFuture);