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

@@ -51,7 +51,7 @@
  * session. Used with {@link com.google.api.gax.retrying.ScheduledRetryingExecutor}.
  *
  * <p>Execution follows the standard attempt template with pre-attempt recovery handling. When the
- * previous attempt failed with a Category 2 (recoverable) error or missing status header, {@code
+ * previous attempt failed with a recoverable error or missing status header, {@code
  * prepareAttempt} queries session status, realigns the buffer window, tops up from the stream, and
  * dispatches the chunk upload request. The callable never blocks on {@code .get()}; results and
  * cancellations propagate asynchronously.
@@ -68,6 +68,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
   private final RewindableStreamBuffer buffer;
   private final String uploadUrl;
   private final ApiCallContext originalCallContext;
+  private final @Nullable UploadProgressTracker progressTracker;
 
   private volatile ChunkUploadRequest currentRequest;
   private volatile UploadCommand currentCommand;
@@ -84,7 +85,8 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
       String uploadUrl,
       ChunkUploadRequest request,
       ApiCallContext callContext,
-      UploadCommand command) {
+      UploadCommand command,
+      @Nullable UploadProgressTracker progressTracker) {
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
     this.queryStatusCallable =
@@ -94,6 +96,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 = progressTracker;
   }
 
   void setRetryingFuture(RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture) {
@@ -128,6 +131,9 @@ private void prepareAttempt(
       SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
       ApiCallContext attemptContext,
       RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    if (progressTracker != null) {
+      progressTracker.onRecovering(lastFailure);
+    }
     QueryStatusRequest queryRequest = QueryStatusRequest.create(uploadUrl);
     ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
         queryStatusCallable.futureCall(queryRequest, attemptContext);
@@ -190,14 +196,18 @@ private void handleQuerySuccess(
       failAttempt(
           attemptFuture,
           RewindableStreamBuffer.protocolViolation(
-              "Incomplete query status response did not include a committed offset for upload URL: "
+              "Incomplete query status response did not include a committed offset for upload"
+                  + " URL: "
                   + uploadUrl));
       return;
     }
 
     // Normal path: realign buffer to committedOffset, compact and top up.
     try {
       buffer.realignTo(committedOffset);
+      if (progressTracker != null) {
+        progressTracker.onOffsetReceived(committedOffset);
+      }
     } catch (Throwable e) {
       failAttempt(attemptFuture, e);
       return;
@@ -207,7 +217,8 @@ private void handleQuerySuccess(
     // Preserve upload,finalize for a trailing partial after realignment.
     UploadCommand realignedCommand;
     if (buffer.isFinal()) {
-      realignedCommand = buffer.isEmpty() ? UploadCommand.FINALIZE : UploadCommand.UPLOAD_FINALIZE;
+      realignedCommand =
+          buffer.isEmpty() ? UploadCommand.FINALIZE : UploadCommand.UPLOAD_FINALIZE;
     } else {
       realignedCommand = UploadCommand.UPLOAD;
     }

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;
@@ -92,6 +93,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
   private final ApiCallContext callContext;
   private final ClientContext clientContext;
   private final RetrySettings chunkRetrySettings;
+  private final UploadProgressTracker progressTracker = new UploadProgressTracker();
 
   private volatile @Nullable String uploadSessionUrl;
   private volatile @Nullable RewindableStreamBuffer buffer;
@@ -183,6 +185,7 @@ public void onSuccess(ResumableUploadSession session) {
               }
             }
             uploadSessionUrl = session.getUploadUrl();
+            progressTracker.onStarted(uploadSessionUrl);
             buffer = new RewindableStreamBuffer(payload, chunkSize, uploadSessionUrl);
             transmitChunk(0L);
           }
@@ -198,6 +201,14 @@ public void onFailure(Throwable t) {
         MoreExecutors.directExecutor());
   }
 
+  void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
+    progressTracker.addListener(listener, executor);
+  }
+
+  ResumableUploadStatus getStatus() {
+    return progressTracker.getStatus();
+  }
+
   private static final StatusCode TIMEOUT_STATUS_CODE =
       new StatusCode() {
         @Override
@@ -210,7 +221,6 @@ public StatusCode.Code getCode() {
           return null;
         }
       };
-
   private void onTimeout() {
     synchronized (lock) {
       if (done) {
@@ -261,6 +271,8 @@ void cancel(boolean mayInterruptIfRunning) {
     if (inFlight != null) {
       inFlight.cancel(mayInterruptIfRunning);
     }
+    progressTracker.onFailed(
+        new CancellationException("Upload was cancelled"), uploadSessionUrl);
     closePayload();
   }
 
@@ -285,8 +297,15 @@ private void finish(@Nullable ResponseT response, @Nullable Throwable error) {
     }
     IOException closeError = closePayload();
     if (error == null) {
+      RewindableStreamBuffer currentBuffer = buffer;
+      long totalBytes =
+          currentBuffer != null
+              ? currentBuffer.getBufferBaseOffset() + currentBuffer.getPayloadLength()
+              : 0L;
+      progressTracker.onFinalized(totalBytes);
       result.set(response);
     } else {
+      progressTracker.onFailed(error, uploadSessionUrl);
       if (closeError != null) {
         error.addSuppressed(closeError);
       }
@@ -360,7 +379,8 @@ private void transmitChunk(long currentOffset) {
             url,
             chunkRequest,
             callContext,
-            command);
+            command,
+            progressTracker);
 
     RetryAlgorithm<ChunkUploadResponse<ResponseT>> retryAlgorithm =
         new RetryAlgorithm<>(
@@ -386,6 +406,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

@@ -98,6 +98,16 @@ void setInFlightFuture(ApiFuture<?> future) {
     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);
@@ -120,12 +130,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

@@ -129,7 +129,8 @@ void call_successfulChunk_setsAttemptFuture() throws Exception {
             "https://upload.url/test",
             request,
             callContext,
-            UploadCommand.UPLOAD);
+            UploadCommand.UPLOAD,
+            /* progressTracker= */ null);
 
     callable.setRetryingFuture(mockExternalFuture);
     ChunkUploadResponse<String> callResult = callable.call();
@@ -177,7 +178,8 @@ void call_returnsWithoutBlocking_andPropagatesCancellation() throws Exception {
             "https://upload.url/test",
             request,
             callContext,
-            UploadCommand.UPLOAD);
+            UploadCommand.UPLOAD,
+            /* progressTracker= */ null);
 
     List<Runnable> listeners = new ArrayList<>();
     doAnswer(
@@ -250,7 +252,8 @@ void call_perAttemptDeadline_appliesRpcTimeoutToCallContext() throws Exception {
             "https://upload.url/test",
             request,
             callContext,
-            UploadCommand.UPLOAD);
+            UploadCommand.UPLOAD,
+            /* progressTracker= */ null);
 
     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,
+            /* progressTracker= */ null);
 
     callable.setRetryingFuture(mockExternalFuture);