Skip to content

Commit 04213cd

Browse files
committed
feat(gax): wire progress listener into coordinator and attempt lifecycle
1 parent 7084b3f commit 04213cd

4 files changed

Lines changed: 388 additions & 20 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
@@ -70,6 +70,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
7070
private final String uploadUrl;
7171
private final RewindableStreamBuffer buffer;
7272
private final ApiCallContext callContext;
73+
private final UploadProgressTracker progressTracker;
7374
private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
7475
private volatile @Nullable ApiFuture<?> currentChunkFuture;
7576

@@ -79,14 +80,16 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
7980
String uploadUrl,
8081
InputStream payload,
8182
int chunkSize,
82-
ApiCallContext callContext) {
83+
ApiCallContext callContext,
84+
UploadProgressTracker progressTracker) {
8385
this.uploadChunkCallable =
8486
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
8587
this.queryStatusCallable =
8688
checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
8789
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
8890
checkNotNull(payload, "payload must not be null");
8991
this.callContext = checkNotNull(callContext, "callContext must not be null");
92+
this.progressTracker = checkNotNull(progressTracker, "progressTracker must not be null");
9093
this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl);
9194
}
9295

@@ -136,7 +139,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
136139
return;
137140
}
138141
if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
139-
recover();
142+
recover(null);
140143
} else {
141144
onChunkUploaded(response);
142145
}
@@ -150,7 +153,7 @@ public void onFailure(Throwable t) {
150153
Category category =
151154
ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD);
152155
if (category == Category.RECOVERABLE) {
153-
recover();
156+
recover(t);
154157
} else {
155158
// Category.TRANSIENT errors reaching here have already exhausted their retry budget
156159
// in the underlying RetryingCallable and become fatal per protocol specification.
@@ -164,7 +167,14 @@ public void onFailure(Throwable t) {
164167
}
165168
}
166169

167-
private void recover() {
170+
/**
171+
* Queries the session for the server's committed offset and resumes transmission from it.
172+
*
173+
* @param cause the error that triggered recovery, or null if the server acknowledged the chunk
174+
* without an upload status header
175+
*/
176+
private void recover(@Nullable Throwable cause) {
177+
progressTracker.onRecovering(cause);
168178
try {
169179
// Dispatch the query status call and register the in-flight future for cancellation.
170180
ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
@@ -228,22 +238,29 @@ private void handleQueryResponse(QueryStatusResponse<ResponseT> queryResponse)
228238
"Incomplete query status response did not include a committed offset for upload URL: "
229239
+ uploadUrl);
230240
}
241+
progressTracker.onOffsetReceived(committedOffset);
231242
buffer.realignTo(committedOffset);
232243
dispatchCurrentChunk();
233244
}
234245

235246
private void onChunkUploaded(ChunkUploadResponse<ResponseT> response) {
236-
long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
237-
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
238-
result.set(response.getResponse());
239-
} else if (buffer.isFinal()) {
240-
result.setException(
241-
new IllegalStateException(
242-
"Upload stream ended and final chunk was transmitted, but server returned"
243-
+ " incomplete status for upload URL: "
244-
+ uploadUrl));
245-
} else {
246-
chunkExecutor.execute(() -> transmitChunk(nextOffset));
247+
try {
248+
long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
249+
if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
250+
progressTracker.onChunkUploaded(nextOffset);
251+
result.set(response.getResponse());
252+
} else if (buffer.isFinal()) {
253+
result.setException(
254+
new IllegalStateException(
255+
"Upload stream ended and final chunk was transmitted, but server returned"
256+
+ " incomplete status for upload URL: "
257+
+ uploadUrl));
258+
} else {
259+
progressTracker.onChunkUploaded(nextOffset);
260+
chunkExecutor.execute(() -> transmitChunk(nextOffset));
261+
}
262+
} catch (Throwable t) {
263+
result.setException(t);
247264
}
248265
}
249266

‎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: 37 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,10 +78,14 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
7878
private final ResumableUploadCallSettings settings;
7979
private final ApiCallContext callContext;
8080
private final ScheduledExecutorService executor;
81+
private final UploadProgressTracker progressTracker = new UploadProgressTracker();
8182
private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
8283

8384
private volatile @Nullable String uploadSessionUrl;
8485

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

@@ -147,18 +151,25 @@ private void start() {
147151
new ApiFutureCallback<ResumableUploadSession>() {
148152
@Override
149153
public void onSuccess(ResumableUploadSession session) {
154+
synchronized (lock) {
155+
if (done) {
156+
return;
157+
}
158+
uploadSessionUrl = session.getUploadUrl();
159+
}
160+
progressTracker.onStarted(uploadSessionUrl);
150161
if (resultFuture.isDone()) {
151162
return;
152163
}
153-
uploadSessionUrl = session.getUploadUrl();
154164
ResumableUploadChunkCoordinator<ResponseT> coordinator =
155165
new ResumableUploadChunkCoordinator<>(
156166
uploadChunkCallable,
157167
queryStatusCallable,
158168
uploadSessionUrl,
159169
payload,
160170
settings.getChunkSize(),
161-
callContext);
171+
callContext,
172+
progressTracker);
162173
ApiFuture<ResponseT> uploadFuture;
163174
try {
164175
uploadFuture = coordinator.start();
@@ -168,7 +179,7 @@ public void onSuccess(ResumableUploadSession session) {
168179
}
169180
boolean alreadyDone = false;
170181
synchronized (lock) {
171-
if (resultFuture.isDone()) {
182+
if (done) {
172183
alreadyDone = true;
173184
} else {
174185
inFlightFuture = uploadFuture;
@@ -219,24 +230,31 @@ private void onTimeout() {
219230

220231
private void succeed(@Nullable ResponseT result) {
221232
synchronized (lock) {
233+
if (done) {
234+
return;
235+
}
236+
done = true;
222237
inFlightFuture = null;
223238
}
239+
progressTracker.onFinalized(progressTracker.getStatus().getBytesUploaded());
224240
closePayload();
225241
resultFuture.set(result);
226242
}
227243

228244
private void fail(Throwable t) {
229245
ApiFuture<?> inFlight;
230246
synchronized (lock) {
231-
if (resultFuture.isDone()) {
247+
if (done) {
232248
return;
233249
}
250+
done = true;
234251
inFlight = this.inFlightFuture;
235252
this.inFlightFuture = null;
236253
}
237254
if (inFlight != null) {
238255
inFlight.cancel(true);
239256
}
257+
progressTracker.onFailed(t, uploadSessionUrl);
240258
closePayload();
241259
resultFuture.setException(t);
242260
}
@@ -254,6 +272,16 @@ private void closePayload() {
254272
return uploadSessionUrl;
255273
}
256274

275+
@Override
276+
public void addProgressListener(ResumableUploadProgressListener listener, Executor executor) {
277+
progressTracker.addListener(listener, executor);
278+
}
279+
280+
@Override
281+
public ResumableUploadProgress getStatus() {
282+
return progressTracker.getStatus();
283+
}
284+
257285
@Override
258286
public void addListener(Runnable listener, Executor executor) {
259287
resultFuture.addListener(listener, executor);
@@ -264,13 +292,18 @@ public boolean cancel(boolean mayInterruptIfRunning) {
264292
boolean cancelled;
265293
ApiFuture<?> inFlight;
266294
synchronized (lock) {
295+
if (done) {
296+
return false;
297+
}
298+
done = true;
267299
cancelled = resultFuture.cancel(mayInterruptIfRunning);
268300
inFlight = this.inFlightFuture;
269301
this.inFlightFuture = null;
270302
}
271303
if (inFlight != null) {
272304
inFlight.cancel(mayInterruptIfRunning);
273305
}
306+
progressTracker.onFailed(new CancellationException("Upload was cancelled"), uploadSessionUrl);
274307
closePayload();
275308
return cancelled;
276309
}

0 commit comments

Comments
 (0)