Skip to content

Commit 9c9b1b5

Browse files
committed
feat(gax): wire progress listener into coordinator and attempt lifecycle
1 parent b9d6cb4 commit 9c9b1b5

4 files changed

Lines changed: 390 additions & 19 deletions

File tree

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

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
7878
private final long deadlineNanos;
7979
private final ApiCallContext callContext;
8080
private final ClientContext clientContext;
81+
private final UploadProgressTracker progressTracker;
8182
private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
8283
private volatile @Nullable ApiFuture<?> currentChunkFuture;
8384

@@ -87,7 +88,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
8788
String uploadUrl,
8889
ResumableUploadCallSettings settings,
8990
long deadlineNanos,
90-
ClientContext clientContext) {
91+
ClientContext clientContext,
92+
UploadProgressTracker progressTracker) {
9193
this.uploadChunkCallable =
9294
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
9395
this.queryStatusCallable =
@@ -98,6 +100,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
98100
this.deadlineNanos = deadlineNanos;
99101
this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
100102
this.callContext = clientContext.getDefaultCallContext();
103+
this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
101104
this.buffer = new RewindableStreamBuffer(payload, settings.getChunkSize(), uploadUrl);
102105
}
103106

@@ -145,7 +148,7 @@ private void dispatchCurrentChunk() {
145148
@Override
146149
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
147150
if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
148-
recover();
151+
recover(null);
149152
} else {
150153
onChunkUploaded(response);
151154
}
@@ -159,7 +162,7 @@ public void onFailure(Throwable t) {
159162
Category category =
160163
ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD);
161164
if (category == Category.RECOVERABLE) {
162-
recover();
165+
recover(t);
163166
} else {
164167
// Category.TRANSIENT errors reaching here have already exhausted their retry budget
165168
// in the underlying RetryingCallable and become fatal per protocol specification.
@@ -173,10 +176,17 @@ public void onFailure(Throwable t) {
173176
}
174177
}
175178

176-
private void recover() {
179+
/**
180+
* Queries the session for the server's committed offset and resumes transmission from it.
181+
*
182+
* @param cause the error that triggered recovery, or null if the server acknowledged the chunk
183+
* without an upload status header
184+
*/
185+
private void recover(@Nullable Throwable cause) {
177186
if (result.isDone()) {
178187
return;
179188
}
189+
progressTracker.onRecovering(cause);
180190
try {
181191
ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
182192
queryStatusCallable.futureCall(QueryStatusRequest.create(uploadUrl), queryCallContext());
@@ -230,6 +240,7 @@ private void handleQueryResponse(QueryStatusResponse<ResponseT> queryResponse) t
230240
"Incomplete query status response did not include a committed offset for upload URL: "
231241
+ uploadUrl);
232242
}
243+
progressTracker.onOffsetReceived(committedOffset);
233244
buffer.realignTo(committedOffset);
234245
dispatchCurrentChunk();
235246
}
@@ -238,17 +249,23 @@ private void onChunkUploaded(ChunkUploadResponse<ResponseT> response) {
238249
if (result.isDone()) {
239250
return;
240251
}
241-
long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
242-
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
243-
result.set(response.getResponse());
244-
} else if (buffer.isFinal()) {
245-
result.setException(
246-
new IllegalStateException(
247-
"Upload stream ended and final chunk was transmitted, but server returned"
248-
+ " incomplete status for upload URL: "
249-
+ uploadUrl));
250-
} else {
251-
chunkExecutor.execute(() -> transmitChunk(nextOffset));
252+
try {
253+
long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
254+
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
255+
progressTracker.onChunkUploaded(nextOffset);
256+
result.set(response.getResponse());
257+
} else if (buffer.isFinal()) {
258+
result.setException(
259+
new IllegalStateException(
260+
"Upload stream ended and final chunk was transmitted, but server returned"
261+
+ " incomplete status for upload URL: "
262+
+ uploadUrl));
263+
} else {
264+
progressTracker.onChunkUploaded(nextOffset);
265+
chunkExecutor.execute(() -> transmitChunk(nextOffset));
266+
}
267+
} catch (Throwable t) {
268+
result.setException(t);
252269
}
253270
}
254271

‎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: 39 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -76,11 +76,15 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
7676
private final ResumableUploadCallSettings settings;
7777
private final ClientContext clientContext;
7878
private final ScheduledExecutorService executor;
79+
private final UploadProgressTracker progressTracker = new UploadProgressTracker();
7980
private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
8081

8182
private volatile @Nullable String uploadSessionUrl;
8283
private volatile long deadlineNanos;
8384

85+
@GuardedBy("lock")
86+
private boolean done;
87+
8488
@GuardedBy("lock")
8589
private @Nullable ApiFuture<?> inFlightFuture;
8690

@@ -149,10 +153,16 @@ private void start() {
149153
new ApiFutureCallback<ResumableUploadSession>() {
150154
@Override
151155
public void onSuccess(ResumableUploadSession session) {
156+
synchronized (lock) {
157+
if (done) {
158+
return;
159+
}
160+
uploadSessionUrl = session.getUploadUrl();
161+
}
162+
progressTracker.onStarted(uploadSessionUrl);
152163
if (resultFuture.isDone()) {
153164
return;
154165
}
155-
uploadSessionUrl = session.getUploadUrl();
156166
ResumableUploadChunkCoordinator<ResponseT> coordinator =
157167
new ResumableUploadChunkCoordinator<>(
158168
uploadChunkCallable,
@@ -161,7 +171,8 @@ public void onSuccess(ResumableUploadSession session) {
161171
payload,
162172
settings,
163173
deadlineNanos,
164-
clientContext);
174+
clientContext,
175+
progressTracker);
165176
ApiFuture<ResponseT> uploadFuture;
166177
try {
167178
uploadFuture = coordinator.start();
@@ -171,7 +182,7 @@ public void onSuccess(ResumableUploadSession session) {
171182
}
172183
boolean alreadyDone = false;
173184
synchronized (lock) {
174-
if (resultFuture.isDone()) {
185+
if (done) {
175186
alreadyDone = true;
176187
} else {
177188
inFlightFuture = uploadFuture;
@@ -224,13 +235,18 @@ private void onTimeout() {
224235
private void succeed(@Nullable ResponseT result) {
225236
ScheduledFuture<?> timeout;
226237
synchronized (lock) {
238+
if (done) {
239+
return;
240+
}
241+
done = true;
227242
inFlightFuture = null;
228243
timeout = this.timeoutFuture;
229244
this.timeoutFuture = null;
230245
}
231246
if (timeout != null) {
232247
timeout.cancel(false);
233248
}
249+
progressTracker.onFinalized(progressTracker.getStatus().getBytesUploaded());
234250
closePayload();
235251
resultFuture.set(result);
236252
}
@@ -239,6 +255,10 @@ private void fail(Throwable t) {
239255
ScheduledFuture<?> timeout;
240256
ApiFuture<?> inFlight;
241257
synchronized (lock) {
258+
if (done) {
259+
return;
260+
}
261+
done = true;
242262
inFlight = this.inFlightFuture;
243263
this.inFlightFuture = null;
244264
timeout = this.timeoutFuture;
@@ -250,6 +270,7 @@ private void fail(Throwable t) {
250270
if (inFlight != null) {
251271
inFlight.cancel(true);
252272
}
273+
progressTracker.onFailed(t, uploadSessionUrl);
253274
closePayload();
254275
resultFuture.setException(t);
255276
}
@@ -267,6 +288,16 @@ private void closePayload() {
267288
return uploadSessionUrl;
268289
}
269290

291+
@Override
292+
public void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
293+
progressTracker.addListener(listener, executor);
294+
}
295+
296+
@Override
297+
public ResumableUploadProgress getStatus() {
298+
return progressTracker.getStatus();
299+
}
300+
270301
@Override
271302
public void addListener(Runnable listener, Executor executor) {
272303
resultFuture.addListener(listener, executor);
@@ -278,6 +309,10 @@ public boolean cancel(boolean mayInterruptIfRunning) {
278309
ApiFuture<?> inFlight;
279310
ScheduledFuture<?> timeout;
280311
synchronized (lock) {
312+
if (done) {
313+
return false;
314+
}
315+
done = true;
281316
cancelled = resultFuture.cancel(mayInterruptIfRunning);
282317
inFlight = this.inFlightFuture;
283318
this.inFlightFuture = null;
@@ -290,6 +325,7 @@ public boolean cancel(boolean mayInterruptIfRunning) {
290325
if (inFlight != null) {
291326
inFlight.cancel(mayInterruptIfRunning);
292327
}
328+
progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
293329
closePayload();
294330
return cancelled;
295331
}

0 commit comments

Comments
 (0)