Skip to content

Commit 49e3dcf

Browse files
committed
feat(gax): wire progress listener into coordinator and attempt lifecycle
1 parent 32dcea1 commit 49e3dcf

8 files changed

Lines changed: 373 additions & 10 deletions

File tree

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
7272
private final RewindableStreamBuffer buffer;
7373
private final String uploadUrl;
7474
private final ApiCallContext originalCallContext;
75+
private final UploadProgressTracker progressTracker;
7576
private final long deadlineNanos;
7677
private final ApiClock clock;
7778

@@ -90,7 +91,8 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
9091
String uploadUrl,
9192
ChunkUploadRequest request,
9293
ApiCallContext callContext,
93-
ResumableUploadCommand command) {
94+
ResumableUploadCommand command,
95+
UploadProgressTracker progressTracker) {
9496
this(
9597
uploadChunkCallable,
9698
queryStatusCallable,
@@ -99,6 +101,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
99101
request,
100102
callContext,
101103
command,
104+
progressTracker,
102105
Long.MAX_VALUE,
103106
NanoClock.getDefaultClock());
104107
}
@@ -111,6 +114,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
111114
ChunkUploadRequest request,
112115
ApiCallContext callContext,
113116
ResumableUploadCommand command,
117+
UploadProgressTracker progressTracker,
114118
long deadlineNanos,
115119
ApiClock clock) {
116120
this.uploadChunkCallable =
@@ -122,6 +126,7 @@ class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<Re
122126
this.currentRequest = checkNotNull(request, "request must not be null");
123127
this.originalCallContext = checkNotNull(callContext, "callContext must not be null");
124128
this.currentCommand = checkNotNull(command, "command must not be null");
129+
this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
125130
this.deadlineNanos = deadlineNanos;
126131
this.clock = checkNotNull(clock, "clock must not be null");
127132
}
@@ -164,6 +169,8 @@ private void prepareAttempt(
164169
ApiCallContext attemptContext,
165170
RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
166171
this.currentCommand = ResumableUploadCommand.QUERY;
172+
progressTracker.onRecovering(lastFailure);
173+
167174
// Per GAX-R7: query uses sensible unary defaults trimmed to the remaining global deadline.
168175
long remainingNanos =
169176
deadlineNanos == Long.MAX_VALUE
@@ -259,6 +266,7 @@ private void handleQuerySuccess(
259266
// Normal path: realign buffer to committedOffset, compact and top up.
260267
try {
261268
buffer.realignTo(committedOffset);
269+
progressTracker.onOffsetReceived(committedOffset);
262270
} catch (Throwable e) {
263271
failAttempt(attemptFuture, e);
264272
return;

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
8787
private final long deadlineNanos;
8888
private final ApiCallContext callContext;
8989
private final ClientContext clientContext;
90+
private final UploadProgressTracker progressTracker;
9091
private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
9192
private volatile @Nullable ApiFuture<?> currentChunkFuture;
9293

@@ -98,7 +99,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
9899
ResumableUploadCallSettings settings,
99100
long deadlineNanos,
100101
ApiCallContext callContext,
101-
ClientContext clientContext) {
102+
ClientContext clientContext,
103+
UploadProgressTracker progressTracker) {
102104
this.uploadChunkCallable =
103105
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
104106
this.queryStatusCallable =
@@ -109,6 +111,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
109111
this.deadlineNanos = deadlineNanos;
110112
this.callContext = checkNotNull(callContext, "callContext must not be null");
111113
this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
114+
this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
112115
this.buffer = new RewindableStreamBuffer(payload, settings.getChunkSize(), uploadUrl);
113116
}
114117

@@ -200,6 +203,7 @@ private void transmitSingleChunk(long currentOffset) {
200203
chunkRequest,
201204
chunkCallContext,
202205
command,
206+
progressTracker,
203207
deadlineNanos,
204208
clientContext.getClock());
205209

@@ -224,6 +228,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
224228
return;
225229
}
226230
long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
231+
progressTracker.onChunkUploaded(nextOffset);
227232
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
228233
result.set(response.getResponse());
229234
} else if (buffer.isFinal()) {

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

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131

3232
import com.google.api.core.ApiFuture;
3333
import com.google.api.core.BetaApi;
34+
import java.util.concurrent.Executor;
3435
import org.jspecify.annotations.NullMarked;
3536
import org.jspecify.annotations.Nullable;
3637

@@ -48,4 +49,18 @@ public interface ResumableUploadFuture<ResponseT> extends ApiFuture<ResponseT> {
4849

4950
/** Returns the upload session URL, or {@code null} if session initiation is in progress. */
5051
@Nullable String getUploadSessionUrl();
52+
53+
/**
54+
* Registers a listener to receive progress and state transition notifications for this upload.
55+
*
56+
* <p>A snapshot of the current upload status is dispatched to the listener immediately upon
57+
* subscription on the provided executor. Subsequent status updates are delivered in order.
58+
*
59+
* @param listener callback listener to receive progress notifications
60+
* @param executor executor on which the listener callbacks are dispatched
61+
*/
62+
void addProgressListener(ResumableUploadProgressListener listener, Executor executor);
63+
64+
/** Returns the current progress snapshot of the upload session. */
65+
ResumableUploadProgress getStatus();
5166
}

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

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
7777
private final ApiCallContext callContext;
7878
private final ClientContext clientContext;
7979
private final ScheduledExecutorService executor;
80+
private final UploadProgressTracker progressTracker = new UploadProgressTracker();
8081
private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
8182

8283
private volatile @Nullable String uploadSessionUrl;
@@ -151,6 +152,7 @@ public void onSuccess(ResumableUploadSession session) {
151152
return;
152153
}
153154
uploadSessionUrl = session.getUploadUrl();
155+
progressTracker.onStarted(uploadSessionUrl);
154156
ResumableUploadChunkCoordinator<ResponseT> coordinator =
155157
new ResumableUploadChunkCoordinator<>(
156158
uploadChunkCallable,
@@ -160,7 +162,8 @@ public void onSuccess(ResumableUploadSession session) {
160162
settings,
161163
deadlineNanos,
162164
callContext,
163-
clientContext);
165+
clientContext,
166+
progressTracker);
164167
ApiFuture<ResponseT> uploadFuture;
165168
try {
166169
uploadFuture = coordinator.start();
@@ -224,6 +227,7 @@ private void succeed(@Nullable ResponseT result) {
224227
if (timeout != null) {
225228
timeout.cancel(false);
226229
}
230+
progressTracker.onFinalized(progressTracker.getStatus().getBytesUploaded());
227231
closePayload();
228232
resultFuture.set(result);
229233
}
@@ -243,6 +247,7 @@ private void fail(Throwable t) {
243247
if (inFlight != null) {
244248
inFlight.cancel(true);
245249
}
250+
progressTracker.onFailed(t, uploadSessionUrl);
246251
closePayload();
247252
resultFuture.setException(t);
248253
}
@@ -260,6 +265,16 @@ private void closePayload() {
260265
return uploadSessionUrl;
261266
}
262267

268+
@Override
269+
public void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
270+
progressTracker.addListener(listener, executor);
271+
}
272+
273+
@Override
274+
public ResumableUploadProgress getStatus() {
275+
return progressTracker.getStatus();
276+
}
277+
263278
@Override
264279
public void addListener(Runnable listener, Executor executor) {
265280
resultFuture.addListener(listener, executor);
@@ -283,6 +298,7 @@ public boolean cancel(boolean mayInterruptIfRunning) {
283298
if (inFlight != null) {
284299
inFlight.cancel(mayInterruptIfRunning);
285300
}
301+
progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
286302
closePayload();
287303
return cancelled;
288304
}

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -131,7 +131,8 @@ void call_successfulChunk_setsAttemptFuture() throws Exception {
131131
"https://upload.url/test",
132132
request,
133133
callContext,
134-
ResumableUploadCommand.UPLOAD);
134+
ResumableUploadCommand.UPLOAD,
135+
new UploadProgressTracker());
135136

136137
callable.setRetryingFuture(mockExternalFuture);
137138
ChunkUploadResponse<String> callResult = callable.call();
@@ -179,7 +180,8 @@ void call_returnsWithoutBlocking_andPropagatesCancellation() throws Exception {
179180
"https://upload.url/test",
180181
request,
181182
callContext,
182-
ResumableUploadCommand.UPLOAD);
183+
ResumableUploadCommand.UPLOAD,
184+
new UploadProgressTracker());
183185

184186
List<Runnable> listeners = new ArrayList<>();
185187
doAnswer(
@@ -253,7 +255,8 @@ void call_perAttemptDeadline_appliesRpcTimeoutToCallContext() throws Exception {
253255
"https://upload.url/test",
254256
request,
255257
callContext,
256-
ResumableUploadCommand.UPLOAD);
258+
ResumableUploadCommand.UPLOAD,
259+
new UploadProgressTracker());
257260

258261
callable.setRetryingFuture(mockExternalFuture);
259262
callable.call();
@@ -294,7 +297,8 @@ void call_nonBlockingExecution_callingThreadMakesImmediateProgress() throws Exce
294297
"https://upload.url/test",
295298
request,
296299
callContext,
297-
ResumableUploadCommand.UPLOAD);
300+
ResumableUploadCommand.UPLOAD,
301+
new UploadProgressTracker());
298302

299303
callable.setRetryingFuture(mockExternalFuture);
300304

0 commit comments

Comments
 (0)