Skip to content

Commit d186a68

Browse files
whoweswhowes
committed
feat(gax): retry chunk upload on transient errors
Introduce ScheduledRetryingExecutor-based chunk retry orchestration in ResumableUploadChunkCoordinator using ExponentialRetryAlgorithm directly, giving each chunk its own fresh attempt budget while global session duration is governed independently. Add ChunkAttemptCallable following the standard six-step attempt template with an attempt preparation seam for subsequent recovery handling, and cancellation propagation to the underlying HTTP future.
1 parent 8f03a72 commit d186a68

7 files changed

Lines changed: 920 additions & 38 deletions

File tree

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

Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
* * Neither the name of Google LLC nor the names of its
15+
* contributors may be used to endorse or promote products derived from
16+
* this software without specific prior written permission.
17+
*
18+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
*/
30+
package com.google.api.gax.rpc;
31+
32+
import static com.google.common.base.Preconditions.checkNotNull;
33+
34+
import com.google.api.core.ApiFuture;
35+
import com.google.api.core.ApiFutureCallback;
36+
import com.google.api.core.ApiFutures;
37+
import com.google.api.gax.resumable.ChunkUploadRequest;
38+
import com.google.api.gax.resumable.ChunkUploadResponse;
39+
import com.google.api.gax.retrying.NonCancellableFuture;
40+
import com.google.api.gax.retrying.RetryingFuture;
41+
import com.google.common.util.concurrent.MoreExecutors;
42+
import java.time.Duration;
43+
import java.util.concurrent.Callable;
44+
import org.jspecify.annotations.NullMarked;
45+
import org.jspecify.annotations.Nullable;
46+
47+
/**
48+
* A {@link Callable} representing an attempt to transmit a single chunk in a resumable upload
49+
* session. Used with {@link com.google.api.gax.retrying.ScheduledRetryingExecutor}.
50+
*
51+
* <p>Execution follows the standard six-step attempt template with an initial attempt preparation
52+
* seam. The callable never blocks on {@code .get()}; results and cancellations propagate
53+
* asynchronously.
54+
*
55+
* @param <ResponseT> the type of the final response message once the upload completes
56+
*/
57+
@NullMarked
58+
class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<ResponseT>> {
59+
60+
private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
61+
uploadChunkCallable;
62+
private final ChunkUploadRequest request;
63+
private final ApiCallContext originalCallContext;
64+
private final UploadCommand command;
65+
66+
private volatile @Nullable RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture;
67+
private volatile @Nullable Throwable lastFailure;
68+
69+
ChunkAttemptCallable(
70+
UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
71+
ChunkUploadRequest request,
72+
ApiCallContext callContext,
73+
UploadCommand command) {
74+
this.uploadChunkCallable =
75+
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
76+
this.request = checkNotNull(request, "request must not be null");
77+
this.originalCallContext = checkNotNull(callContext, "callContext must not be null");
78+
this.command = checkNotNull(command, "command must not be null");
79+
}
80+
81+
void setRetryingFuture(RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture) {
82+
this.retryingFuture = checkNotNull(retryingFuture, "retryingFuture must not be null");
83+
}
84+
85+
/**
86+
* Pre-attempt hook called before each transmission attempt.
87+
*
88+
* <p>In this phase, Category 2 (recoverable) errors throw to fail fast. Subsequent phases expand
89+
* this seam into the query status -> realign buffer -> top up recovery sequence.
90+
*/
91+
void prepareAttempt() {
92+
if (lastFailure != null) {
93+
UploadErrorCategory category = UploadErrorClassifier.classify(lastFailure, command);
94+
if (category == UploadErrorCategory.RECOVERABLE) {
95+
throw new UnsupportedOperationException(
96+
"Category 2 (recoverable) error recovery is not yet implemented", lastFailure);
97+
}
98+
}
99+
}
100+
101+
@Override
102+
public @Nullable ChunkUploadResponse<ResponseT> call() {
103+
RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture =
104+
checkNotNull(retryingFuture, "retryingFuture must be set before call()");
105+
ApiCallContext attemptContext = originalCallContext;
106+
107+
try {
108+
// Seam for recoverable error handling.
109+
prepareAttempt();
110+
111+
// Set the RPC timeout if caller did not provide their own.
112+
Duration rpcTimeout = currentRetryingFuture.getAttemptSettings().getRpcTimeoutDuration();
113+
if (!rpcTimeout.isZero() && attemptContext.getTimeoutDuration() == null) {
114+
attemptContext = attemptContext.withTimeoutDuration(rpcTimeout);
115+
}
116+
117+
// Placeholder non-cancellable future.
118+
currentRetryingFuture.setAttemptFuture(
119+
new NonCancellableFuture<ChunkUploadResponse<ResponseT>>());
120+
121+
// Early exit if retryingFuture was already cancelled or completed.
122+
if (currentRetryingFuture.isDone()) {
123+
return null;
124+
}
125+
126+
// Dispatch chunk upload and wire cancellation propagation and error
127+
// tracking.
128+
attemptContext
129+
.getTracer()
130+
.attemptStarted(
131+
request, currentRetryingFuture.getAttemptSettings().getOverallAttemptCount());
132+
133+
ApiFuture<ChunkUploadResponse<ResponseT>> internalFuture =
134+
uploadChunkCallable.futureCall(request, attemptContext);
135+
136+
// Propagate cancellation to the feeder future immediately.
137+
currentRetryingFuture.addListener(
138+
() -> {
139+
if (currentRetryingFuture.isCancelled()) {
140+
internalFuture.cancel(true);
141+
}
142+
},
143+
MoreExecutors.directExecutor());
144+
145+
ApiFutures.addCallback(
146+
internalFuture,
147+
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
148+
@Override
149+
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
150+
lastFailure = null;
151+
}
152+
153+
@Override
154+
public void onFailure(Throwable t) {
155+
lastFailure = t;
156+
}
157+
},
158+
MoreExecutors.directExecutor());
159+
160+
currentRetryingFuture.setAttemptFuture(internalFuture);
161+
} catch (Throwable e) {
162+
lastFailure = e;
163+
currentRetryingFuture.setAttemptFuture(
164+
ApiFutures.<ChunkUploadResponse<ResponseT>>immediateFailedFuture(e));
165+
}
166+
167+
return null;
168+
}
169+
}

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ public ResumableUploadCallableImpl(
6868
this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
6969
}
7070

71+
7172
@Override
7273
public ResumableUploadFuture<ResponseT> futureCall(
7374
RequestT request,
@@ -91,7 +92,8 @@ public ResumableUploadFuture<ResponseT> futureCall(
9192
client.uploadChunkCallable(),
9293
payload,
9394
effectiveSettings,
94-
clientContext.getDefaultCallContext());
95+
clientContext.getDefaultCallContext(),
96+
clientContext);
9597
}
9698

9799
@Override

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

Lines changed: 78 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,17 @@
3939
import com.google.api.gax.resumable.ChunkUploadRequest;
4040
import com.google.api.gax.resumable.ChunkUploadResponse;
4141
import com.google.api.gax.resumable.ResumableUploadSession;
42+
import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
43+
import com.google.api.gax.retrying.RetryAlgorithm;
44+
import com.google.api.gax.retrying.RetrySettings;
45+
import com.google.api.gax.retrying.RetryingFuture;
46+
import com.google.api.gax.retrying.ScheduledRetryingExecutor;
4247
import com.google.common.io.ByteStreams;
4348
import com.google.common.util.concurrent.MoreExecutors;
4449
import com.google.errorprone.annotations.concurrent.GuardedBy;
4550
import java.io.IOException;
4651
import java.io.InputStream;
52+
import java.time.Duration;
4753
import java.util.Arrays;
4854
import java.util.concurrent.CancellationException;
4955
import org.jspecify.annotations.NullMarked;
@@ -57,6 +63,18 @@
5763
@NullMarked
5864
final class ResumableUploadChunkCoordinator<ResponseT> {
5965

66+
static final RetrySettings DEFAULT_CHUNK_RETRY_SETTINGS =
67+
RetrySettings.newBuilder()
68+
.setInitialRetryDelayDuration(Duration.ofMillis(100))
69+
.setRetryDelayMultiplier(1.3)
70+
.setMaxRetryDelayDuration(Duration.ofMinutes(1))
71+
.setInitialRpcTimeoutDuration(Duration.ofSeconds(30))
72+
.setRpcTimeoutMultiplier(1.0)
73+
.setMaxRpcTimeoutDuration(Duration.ofSeconds(30))
74+
.setTotalTimeoutDuration(Duration.ofMinutes(5))
75+
.setMaxAttempts(5)
76+
.build();
77+
6078
private static final byte[] EMPTY_PAYLOAD = new byte[0];
6179

6280
private final Object lock = new Object();
@@ -69,6 +87,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
6987
private final byte[] buffer;
7088
private final int chunkSize;
7189
private final ApiCallContext callContext;
90+
private final ClientContext clientContext;
91+
private final RetrySettings chunkRetrySettings;
7292

7393
private volatile @Nullable String uploadSessionUrl;
7494

@@ -87,7 +107,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
87107
UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
88108
InputStream payload,
89109
ResumableUploadCallSettings settings,
90-
ApiCallContext callContext) {
110+
ApiCallContext callContext,
111+
ClientContext clientContext) {
91112
this.result = checkNotNull(result, "result must not be null");
92113
this.startFuture = checkNotNull(startFuture, "startFuture must not be null");
93114
this.uploadChunkCallable =
@@ -97,6 +118,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
97118
checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0");
98119
this.chunkSize = settings.getChunkSize();
99120
this.callContext = checkNotNull(callContext, "callContext must not be null");
121+
this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
122+
this.chunkRetrySettings = DEFAULT_CHUNK_RETRY_SETTINGS;
100123
this.buffer = new byte[chunkSize];
101124
synchronized (lock) {
102125
this.inFlightFuture = startFuture;
@@ -236,47 +259,66 @@ private void transmitChunk(long currentOffset) {
236259
.setFinal(isFinal)
237260
.build();
238261

239-
long chunkLength = chunkPayload.length;
240-
try {
241-
ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
242-
uploadChunkCallable.futureCall(chunkRequest, callContext);
243-
setInFlightFuture(chunkFuture);
262+
UploadCommand command = isFinal ? UploadCommand.UPLOAD_FINALIZE : UploadCommand.UPLOAD;
263+
ChunkAttemptCallable<ResponseT> attemptCallable =
264+
new ChunkAttemptCallable<>(
265+
uploadChunkCallable, chunkRequest, callContext, command);
244266

245-
ApiFutures.addCallback(
246-
chunkFuture,
247-
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
248-
@Override
249-
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
250-
synchronized (lock) {
251-
if (done) {
252-
return;
253-
}
254-
}
255-
long nextOffset = currentOffset + chunkLength;
256-
if (response.isComplete()) {
257-
finish(response.getResponse(), null);
258-
} else if (isFinal) {
259-
finish(
260-
null,
261-
new IllegalStateException(
262-
"Upload stream ended and final chunk was transmitted, but server returned"
263-
+ " incomplete status"));
264-
} else {
265-
transmitChunk(nextOffset);
266-
}
267-
}
267+
RetryAlgorithm<ChunkUploadResponse<ResponseT>> retryAlgorithm =
268+
new RetryAlgorithm<>(
269+
new UploadResultRetryAlgorithm<>(command),
270+
new ExponentialRetryAlgorithm(chunkRetrySettings, clientContext.getClock()));
271+
272+
ScheduledRetryingExecutor<ChunkUploadResponse<ResponseT>> retryingExecutor =
273+
new ScheduledRetryingExecutor<>(retryAlgorithm, clientContext.getExecutor());
268274

269-
@Override
270-
public void onFailure(Throwable t) {
271-
if (t instanceof CancellationException) {
275+
RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture =
276+
retryingExecutor.createFuture(attemptCallable, callContext);
277+
attemptCallable.setRetryingFuture(retryingFuture);
278+
setInFlightFuture(retryingFuture);
279+
280+
long chunkLength = chunkPayload.length;
281+
ApiFutures.addCallback(
282+
retryingFuture,
283+
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
284+
@Override
285+
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
286+
synchronized (lock) {
287+
if (done) {
272288
return;
273289
}
274-
finish(null, t);
275290
}
276-
},
277-
MoreExecutors.directExecutor());
291+
long nextOffset = currentOffset + chunkLength;
292+
if (response.isComplete()) {
293+
finish(response.getResponse(), null);
294+
} else if (isFinal) {
295+
finish(
296+
null,
297+
new IllegalStateException(
298+
"Upload stream ended and final chunk was transmitted, but server returned"
299+
+ " incomplete status for upload URL: "
300+
+ url));
301+
} else {
302+
transmitChunk(nextOffset);
303+
}
304+
}
305+
306+
@Override
307+
public void onFailure(Throwable t) {
308+
if (t instanceof CancellationException) {
309+
return;
310+
}
311+
finish(null, t);
312+
}
313+
},
314+
MoreExecutors.directExecutor());
315+
316+
try {
317+
attemptCallable.call();
278318
} catch (Throwable t) {
279-
finish(null, t);
319+
if (!retryingFuture.isDone()) {
320+
finish(null, t);
321+
}
280322
}
281323
}
282324
}

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

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,32 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
6161
InputStream payload,
6262
ResumableUploadCallSettings settings,
6363
ApiCallContext callContext) {
64+
return create(
65+
startFuture,
66+
uploadChunkCallable,
67+
payload,
68+
settings,
69+
callContext,
70+
ClientContext.newBuilder().setDefaultCallContext(callContext).build());
71+
}
72+
73+
static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
74+
ApiFuture<ResumableUploadSession> startFuture,
75+
UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
76+
InputStream payload,
77+
ResumableUploadCallSettings settings,
78+
ApiCallContext callContext,
79+
ClientContext clientContext) {
6480
SettableApiFuture<ResponseT> result = SettableApiFuture.create();
6581
ResumableUploadChunkCoordinator<ResponseT> coordinator =
6682
new ResumableUploadChunkCoordinator<>(
67-
result, startFuture, uploadChunkCallable, payload, settings, callContext);
83+
result,
84+
startFuture,
85+
uploadChunkCallable,
86+
payload,
87+
settings,
88+
callContext,
89+
clientContext);
6890
ResumableUploadFutureImpl<ResponseT> handle =
6991
new ResumableUploadFutureImpl<>(result, coordinator);
7092
coordinator.start();

0 commit comments

Comments
 (0)