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

@@ -71,6 +71,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;
 
@@ -89,7 +90,8 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
       String uploadUrl,
       ChunkUploadRequest request,
       ApiCallContext callContext,
-      UploadCommand command) {
+      UploadCommand command,
+      UploadProgressTracker progressTracker) {
     this(
         uploadChunkCallable,
         queryStatusCallable,
@@ -98,6 +100,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
         request,
         callContext,
         command,
+        progressTracker,
         Long.MAX_VALUE,
         NanoClock.getDefaultClock());
   }
@@ -110,6 +113,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
       ChunkUploadRequest request,
       ApiCallContext callContext,
       UploadCommand command,
+      UploadProgressTracker progressTracker,
       long deadlineNanos,
       ApiClock clock) {
     this.uploadChunkCallable =
@@ -121,6 +125,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");
   }
@@ -158,6 +163,8 @@ private void prepareAttempt(
       SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
       ApiCallContext attemptContext,
       RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    progressTracker.onRecovering(lastFailure);
+
     // Per GAX-R7: query uses sensible unary defaults trimmed to the remaining global deadline.
     long remainingNanos =
         deadlineNanos == Long.MAX_VALUE
@@ -254,6 +261,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

@@ -52,6 +52,7 @@
 import java.io.InputStream;
 import java.time.Duration;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
@@ -94,6 +95,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
   private final int chunkSize;
   private final ApiCallContext callContext;
   private final ClientContext clientContext;
+  private final UploadProgressTracker progressTracker = new UploadProgressTracker();
 
   private volatile @Nullable String uploadSessionUrl;
   private volatile @Nullable RewindableStreamBuffer buffer;
@@ -166,6 +168,7 @@ public void onSuccess(ResumableUploadSession session) {
               }
             }
             uploadSessionUrl = session.getUploadUrl();
+            progressTracker.onStarted(uploadSessionUrl);
             buffer = new RewindableStreamBuffer(payload, chunkSize, uploadSessionUrl);
             scheduleNextChunk(0L);
           }
@@ -181,6 +184,14 @@ public void onFailure(Throwable t) {
         MoreExecutors.directExecutor());
   }
 
+  void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
+    progressTracker.addListener(listener, executor);
+  }
+
+  ResumableUploadStatus getStatus() {
+    return progressTracker.getStatus();
+  }
+
   private void onTimeout() {
     synchronized (lock) {
       if (done) {
@@ -233,6 +244,7 @@ void cancel(boolean mayInterruptIfRunning) {
     if (inFlight != null) {
       inFlight.cancel(mayInterruptIfRunning);
     }
+    progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
     closePayload();
   }
 
@@ -257,11 +269,15 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) {
     }
     IOException closeError = closePayload();
     if (error == null) {
+      long totalBytes =
+          buffer != null ? buffer.getBufferBaseOffset() + buffer.getPayloadLength() : 0L;
+      progressTracker.onFinalized(totalBytes);
       result.set(response);
     } else {
       if (closeError != null) {
         error.addSuppressed(closeError);
       }
+      progressTracker.onFailed(error, uploadSessionUrl);
       result.setException(error);
     }
   }
@@ -364,6 +380,7 @@ private void transmitSingleChunk(long currentOffset) {
             chunkRequest,
             chunkCallContext,
             command,
+            progressTracker,
             deadlineNanos,
             clientContext.getClock());
 
@@ -383,6 +400,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
               }
             }
             long nextOffset = streamBuffer.getBufferBaseOffset() + streamBuffer.getPayloadLength();
+            progressTracker.onChunkUploaded(nextOffset);
             if (response.isComplete()) {
               finish(response.getResponse(), null);
             } else if (streamBuffer.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 status snapshot of the upload session. */
+  ResumableUploadStatus getStatus();
 }

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

@@ -93,6 +93,16 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
     return coordinator.getUploadSessionUrl();
   }
 
+  @Override
+  public void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
+    coordinator.addProgressListener(listener, executor);
+  }
+
+  @Override
+  public ResumableUploadStatus getStatus() {
+    return coordinator.getStatus();
+  }
+
   @Override
   public void addListener(Runnable listener, Executor executor) {
     result.addListener(listener, executor);
@@ -115,12 +125,12 @@ public boolean isDone() {
   }
 
   @Override
-  public ResponseT get() throws InterruptedException, ExecutionException {
+  public @Nullable ResponseT get() throws InterruptedException, ExecutionException {
     return result.get();
   }
 
   @Override
-  public ResponseT get(long timeout, TimeUnit unit)
+  public @Nullable ResponseT get(long timeout, TimeUnit unit)
       throws InterruptedException, ExecutionException, TimeoutException {
     return result.get(timeout, unit);
   }

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

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