Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
whowes wants to merge 1 commit into from
Draft

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion ...platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ChunkAttemptCallable.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -89,7 +90,8 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
String uploadUrl,
ChunkUploadRequest request,
ApiCallContext callContext,
ResumableUploadCommand command) {
ResumableUploadCommand command,
UploadProgressTracker progressTracker) {
this(
uploadChunkCallable,
queryStatusCallable,
Expand All @@ -98,6 +100,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
request,
callContext,
command,
progressTracker,
Long.MAX_VALUE,
NanoClock.getDefaultClock());
}
Expand All @@ -110,6 +113,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
ChunkUploadRequest request,
ApiCallContext callContext,
ResumableUploadCommand command,
UploadProgressTracker progressTracker,
long deadlineNanos,
ApiClock clock) {
this.uploadChunkCallable =
Expand All @@ -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");
}
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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;
Expand Down
18 changes: 18 additions & 0 deletions ...va/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -166,6 +168,7 @@ public void onSuccess(ResumableUploadSession session) {
}
}
uploadSessionUrl = session.getUploadUrl();
progressTracker.onStarted(uploadSessionUrl);
buffer = new RewindableStreamBuffer(payload, chunkSize, uploadSessionUrl);
scheduleNextChunk(0L);
}
Expand All @@ -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) {
Expand Down Expand Up @@ -233,6 +244,7 @@ void cancel(boolean mayInterruptIfRunning) {
if (inFlight != null) {
inFlight.cancel(mayInterruptIfRunning);
}
progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
closePayload();
Comment on lines 244 to 248

@gemini-code-assist gemini-code-assist Bot Sep 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

There is a race condition in the cancel method because it does not acquire the lock or set the done flag to true. This leads to several critical issues:

  1. Resource Leak / Continued Execution: If cancel is called while session initiation is in progress, and the session initiation subsequently succeeds, the onSuccess callback will check done under lock. Since done is still false, it will proceed to call progressTracker.onStarted and start transmitting chunks via transmitChunk(0L), causing the upload to continue running in the background despite being cancelled.
  2. Inconsistent Progress Events: The progress tracker could receive an onFailed event from cancel, followed by onStarted, onChunkUploaded, or onFinalized events from the continuing upload, violating the state machine guarantees.
  3. Duplicate Failure Events: If inFlight is not null, calling inFlight.cancel will trigger its failure callback, which eventually calls finish(null, cancellationException). Since done was not set to true, finish will proceed and call progressTracker.onFailed a second time.

To fix this, cancel should acquire lock, check if already done, and set done = true before proceeding.

    synchronized (lock) {
      if (done) {
        return;
      }
      done = true;
    }
    if (inFlight != null) {
      inFlight.cancel(mayInterruptIfRunning);
    }
    progressTracker.onFailed(
        new CancellationException("Upload was cancelled"), uploadSessionUrl);
    closePayload();

}

Expand All @@ -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);
}
}
Expand Down Expand Up @@ -367,6 +383,7 @@ private void transmitSingleChunk(long currentOffset) {
chunkRequest,
chunkCallContext,
command,
progressTracker,
deadlineNanos,
clientContext.getClock());

Expand All @@ -386,6 +403,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()) {
Expand Down
15 changes: 15 additions & 0 deletions ...latform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFuture.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand All @@ -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();
}
14 changes: 12 additions & 2 deletions ...orm-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadFutureImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
}
Expand Down
12 changes: 8 additions & 4 deletions ...form-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ChunkAttemptCallableTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,7 +130,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();
Expand Down Expand Up @@ -178,7 +179,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(
Expand Down Expand Up @@ -251,7 +253,8 @@ void call_perAttemptDeadline_appliesRpcTimeoutToCallContext() throws Exception {
"https://upload.url/test",
request,
callContext,
ResumableUploadCommand.UPLOAD);
ResumableUploadCommand.UPLOAD,
new UploadProgressTracker());

callable.setRetryingFuture(mockExternalFuture);
callable.call();
Expand Down Expand Up @@ -292,7 +295,8 @@ void call_nonBlockingExecution_callingThreadMakesImmediateProgress() throws Exce
"https://upload.url/test",
request,
callContext,
ResumableUploadCommand.UPLOAD);
ResumableUploadCommand.UPLOAD,
new UploadProgressTracker());

callable.setRetryingFuture(mockExternalFuture);

Expand Down
Loading
Loading