Skip to content

Commit 1fe75dc

Browse files
whoweswhowes
committed
feat(gax): add recoverable error query and buffer realignment loop
Introduces ChunkAttemptCallable with a query-and-realign loop when encountering recoverable protocol errors. Queries the server for the committed offset and adjusts the buffer window before resuming chunk transmission.
1 parent 00f7932 commit 1fe75dc

12 files changed

Lines changed: 1097 additions & 75 deletions

‎sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/resumable/ChunkUploadResponse.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@ public static <ResponseT> Builder<ResponseT> newBuilder() {
6969

7070
public static <ResponseT> ChunkUploadResponse<ResponseT> create(
7171
boolean isComplete, @Nullable ResponseT response) {
72-
return create(isComplete, response, null);
72+
return create(isComplete, response, isComplete ? "final" : "active");
7373
}
7474

7575
public static <ResponseT> ChunkUploadResponse<ResponseT> create(

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

Lines changed: 301 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,301 @@
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.core.SettableApiFuture;
38+
import com.google.api.gax.resumable.ChunkUploadRequest;
39+
import com.google.api.gax.resumable.ChunkUploadResponse;
40+
import com.google.api.gax.resumable.QueryStatusRequest;
41+
import com.google.api.gax.resumable.QueryStatusResponse;
42+
import com.google.api.gax.retrying.RetryingFuture;
43+
import com.google.common.util.concurrent.MoreExecutors;
44+
import java.time.Duration;
45+
import java.util.concurrent.Callable;
46+
import org.jspecify.annotations.NullMarked;
47+
import org.jspecify.annotations.Nullable;
48+
49+
/**
50+
* A {@link Callable} representing an attempt to transmit a single chunk in a resumable upload
51+
* session. Used with {@link com.google.api.gax.retrying.ScheduledRetryingExecutor}.
52+
*
53+
* <p>Execution follows the standard attempt template with pre-attempt recovery handling. When the
54+
* previous attempt failed with a Category 2 (recoverable) error or missing status header, {@code
55+
* prepareAttempt} queries session status, realigns the buffer window, tops up from the stream, and
56+
* dispatches the chunk upload request. The callable never blocks on {@code .get()}; results and
57+
* cancellations propagate asynchronously.
58+
*
59+
* @param <ResponseT> the type of the final response message once the upload completes
60+
*/
61+
@NullMarked
62+
class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<ResponseT>> {
63+
64+
private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
65+
uploadChunkCallable;
66+
private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
67+
queryStatusCallable;
68+
private final RewindableStreamBuffer buffer;
69+
private final String uploadUrl;
70+
private final ApiCallContext originalCallContext;
71+
72+
private volatile ChunkUploadRequest currentRequest;
73+
private volatile ResumableUploadCommand currentCommand;
74+
75+
private volatile @Nullable RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture;
76+
private volatile @Nullable ApiFuture<?> inFlightFuture;
77+
private volatile @Nullable Throwable lastFailure;
78+
private volatile @Nullable ChunkUploadResponse<ResponseT> lastResponse;
79+
80+
ChunkAttemptCallable(
81+
UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
82+
UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
83+
RewindableStreamBuffer buffer,
84+
String uploadUrl,
85+
ChunkUploadRequest request,
86+
ApiCallContext callContext,
87+
ResumableUploadCommand command) {
88+
this.uploadChunkCallable =
89+
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
90+
this.queryStatusCallable =
91+
checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
92+
this.buffer = checkNotNull(buffer, "buffer must not be null");
93+
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
94+
this.currentRequest = checkNotNull(request, "request must not be null");
95+
this.originalCallContext = checkNotNull(callContext, "callContext must not be null");
96+
this.currentCommand = checkNotNull(command, "command must not be null");
97+
}
98+
99+
void setRetryingFuture(RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture) {
100+
this.retryingFuture = checkNotNull(retryingFuture, "retryingFuture must not be null");
101+
}
102+
103+
private boolean needsRecovery() {
104+
if (lastFailure != null) {
105+
ResumableUploadErrorClassifier.Category category =
106+
ResumableUploadErrorClassifier.classify(lastFailure, currentCommand);
107+
return category == ResumableUploadErrorClassifier.Category.RECOVERABLE;
108+
}
109+
if (lastResponse != null && lastResponse.getUploadStatus() == null) {
110+
ResumableUploadErrorClassifier.Category category =
111+
ResumableUploadErrorClassifier.classifyMissingStatusHeader(currentCommand);
112+
return category == ResumableUploadErrorClassifier.Category.RECOVERABLE;
113+
}
114+
return false;
115+
}
116+
117+
private void failAttempt(
118+
SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture, Throwable t) {
119+
lastFailure = t;
120+
lastResponse = null;
121+
attemptFuture.setException(t);
122+
}
123+
124+
/**
125+
* Pre-attempt recovery step invoked before transmitting an attempt when the previous attempt
126+
* encountered a Category 2 (recoverable) error or missing status header.
127+
*/
128+
private void prepareAttempt(
129+
SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
130+
ApiCallContext attemptContext,
131+
RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
132+
QueryStatusRequest queryRequest = QueryStatusRequest.create(uploadUrl);
133+
ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
134+
queryStatusCallable.futureCall(queryRequest, attemptContext);
135+
if (queryFuture == null) {
136+
failAttempt(
137+
attemptFuture, new IllegalStateException("queryStatusCallable returned a null future"));
138+
return;
139+
}
140+
this.inFlightFuture = queryFuture;
141+
142+
ApiFutures.addCallback(
143+
queryFuture,
144+
new ApiFutureCallback<QueryStatusResponse<ResponseT>>() {
145+
@Override
146+
public void onSuccess(QueryStatusResponse<ResponseT> queryResponse) {
147+
handleQuerySuccess(queryResponse, attemptFuture, attemptContext, currentRetryingFuture);
148+
}
149+
150+
@Override
151+
public void onFailure(Throwable t) {
152+
failAttempt(attemptFuture, t);
153+
}
154+
},
155+
MoreExecutors.directExecutor());
156+
}
157+
158+
private void handleQuerySuccess(
159+
QueryStatusResponse<ResponseT> queryResponse,
160+
SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
161+
ApiCallContext attemptContext,
162+
RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
163+
if (currentRetryingFuture.isDone()) {
164+
return;
165+
}
166+
167+
if (queryResponse.getUploadStatus() == null) {
168+
failAttempt(
169+
attemptFuture,
170+
UploadErrors.protocolViolation(
171+
"Query status response missing X-Goog-Upload-Status header for upload URL: "
172+
+ uploadUrl));
173+
return;
174+
}
175+
176+
// Server already finalized the upload.
177+
if (queryResponse.isComplete()) {
178+
ChunkUploadResponse<ResponseT> response =
179+
ChunkUploadResponse.create(
180+
true, queryResponse.getResponse(), queryResponse.getUploadStatus());
181+
lastFailure = null;
182+
lastResponse = response;
183+
attemptFuture.set(response);
184+
return;
185+
}
186+
187+
// Incomplete query response with null committed offset violates the protocol invariant.
188+
Long committedOffset = queryResponse.getCommittedOffset();
189+
if (committedOffset == null) {
190+
failAttempt(
191+
attemptFuture,
192+
UploadErrors.protocolViolation(
193+
"Incomplete query status response did not include a committed offset for upload URL: "
194+
+ uploadUrl));
195+
return;
196+
}
197+
198+
// Normal path: realign buffer to committedOffset, compact and top up.
199+
try {
200+
buffer.realignTo(committedOffset);
201+
} catch (Throwable e) {
202+
failAttempt(attemptFuture, e);
203+
return;
204+
}
205+
206+
// Determine the upload command for the realigned buffer.
207+
// Preserve upload,finalize for a trailing partial after realignment.
208+
ResumableUploadCommand realignedCommand;
209+
if (buffer.isFinal()) {
210+
realignedCommand =
211+
buffer.isEmpty()
212+
? ResumableUploadCommand.FINALIZE
213+
: ResumableUploadCommand.UPLOAD_FINALIZE;
214+
} else {
215+
realignedCommand = ResumableUploadCommand.UPLOAD;
216+
}
217+
218+
ChunkUploadRequest realignedRequest =
219+
ChunkUploadRequest.newBuilder()
220+
.setUploadUrl(uploadUrl)
221+
.setPayload(buffer.getBuffer())
222+
.setPayloadLength(buffer.getPayloadLength())
223+
.setOffset(buffer.getBufferBaseOffset())
224+
.setFinal(buffer.isFinal())
225+
.build();
226+
227+
this.currentRequest = realignedRequest;
228+
this.currentCommand = realignedCommand;
229+
230+
dispatchChunkUpload(attemptFuture, attemptContext, currentRetryingFuture);
231+
}
232+
233+
private void dispatchChunkUpload(
234+
SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
235+
ApiCallContext attemptContext,
236+
RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
237+
ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
238+
uploadChunkCallable.futureCall(currentRequest, attemptContext);
239+
this.inFlightFuture = chunkFuture;
240+
241+
ApiFutures.addCallback(
242+
chunkFuture,
243+
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
244+
@Override
245+
public void onSuccess(ChunkUploadResponse<ResponseT> response) {
246+
lastFailure = null;
247+
lastResponse = response;
248+
attemptFuture.set(response);
249+
}
250+
251+
@Override
252+
public void onFailure(Throwable t) {
253+
failAttempt(attemptFuture, t);
254+
}
255+
},
256+
MoreExecutors.directExecutor());
257+
}
258+
259+
@Override
260+
public @Nullable ChunkUploadResponse<ResponseT> call() {
261+
RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture =
262+
checkNotNull(retryingFuture, "retryingFuture must be set before call()");
263+
ApiCallContext attemptContext = originalCallContext;
264+
265+
Duration rpcTimeout = currentRetryingFuture.getAttemptSettings().getRpcTimeoutDuration();
266+
if (!rpcTimeout.isZero() && attemptContext.getTimeoutDuration() == null) {
267+
attemptContext = attemptContext.withTimeoutDuration(rpcTimeout);
268+
}
269+
270+
SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture = SettableApiFuture.create();
271+
currentRetryingFuture.setAttemptFuture(attemptFuture);
272+
273+
if (currentRetryingFuture.isDone()) {
274+
return null;
275+
}
276+
277+
currentRetryingFuture.addListener(
278+
() -> {
279+
if (currentRetryingFuture.isCancelled()) {
280+
ApiFuture<?> inFlight = inFlightFuture;
281+
if (inFlight != null) {
282+
inFlight.cancel(true);
283+
}
284+
attemptFuture.cancel(true);
285+
}
286+
},
287+
MoreExecutors.directExecutor());
288+
289+
try {
290+
if (needsRecovery()) {
291+
prepareAttempt(attemptFuture, attemptContext, currentRetryingFuture);
292+
} else {
293+
dispatchChunkUpload(attemptFuture, attemptContext, currentRetryingFuture);
294+
}
295+
} catch (Throwable t) {
296+
failAttempt(attemptFuture, t);
297+
}
298+
299+
return null;
300+
}
301+
}

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,16 @@
3535
import com.google.api.core.ApiFutures;
3636
import com.google.api.core.BetaApi;
3737
import com.google.api.core.InternalApi;
38+
import com.google.api.gax.resumable.QueryStatusRequest;
39+
import com.google.api.gax.resumable.QueryStatusResponse;
3840
import com.google.api.gax.resumable.ResumableUploadClient;
3941
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.ScheduledRetryingExecutor;
4046
import java.io.InputStream;
47+
import java.time.Duration;
4148
import org.jspecify.annotations.NullMarked;
4249
import org.jspecify.annotations.Nullable;
4350

@@ -54,9 +61,22 @@
5461
public class ResumableUploadCallableImpl<RequestT, ResponseT>
5562
extends ResumableUploadCallable<RequestT, ResponseT> {
5663

64+
static final RetrySettings DEFAULT_QUERY_RETRY_SETTINGS =
65+
RetrySettings.newBuilder()
66+
.setInitialRetryDelayDuration(Duration.ofMillis(100))
67+
.setRetryDelayMultiplier(1.3)
68+
.setMaxRetryDelayDuration(Duration.ofMinutes(1))
69+
.setInitialRpcTimeoutDuration(Duration.ofSeconds(30))
70+
.setRpcTimeoutMultiplier(1.0)
71+
.setMaxRpcTimeoutDuration(Duration.ofSeconds(30))
72+
.setTotalTimeoutDuration(Duration.ofMinutes(5))
73+
.build();
74+
5775
private final ResumableUploadClient<RequestT, ResponseT> client;
5876
private final ResumableUploadCallSettings defaultCallSettings;
5977
private final ClientContext clientContext;
78+
private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
79+
retryingQueryCallable;
6080

6181
public ResumableUploadCallableImpl(
6282
ResumableUploadClient<RequestT, ResponseT> client,
@@ -66,6 +86,17 @@ public ResumableUploadCallableImpl(
6686
this.defaultCallSettings =
6787
checkNotNull(defaultCallSettings, "defaultCallSettings must not be null");
6888
this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
89+
90+
RetryAlgorithm<QueryStatusResponse<ResponseT>> queryRetryAlgorithm =
91+
new RetryAlgorithm<>(
92+
new ResumableUploadResultRetryAlgorithm<>(ResumableUploadCommand.QUERY),
93+
new ExponentialRetryAlgorithm(DEFAULT_QUERY_RETRY_SETTINGS, clientContext.getClock()));
94+
95+
this.retryingQueryCallable =
96+
new RetryingCallable<>(
97+
clientContext.getDefaultCallContext(),
98+
checkNotNull(client.queryStatusCallable(), "queryStatusCallable must not be null"),
99+
new ScheduledRetryingExecutor<>(queryRetryAlgorithm, clientContext.getExecutor()));
69100
}
70101

71102
@Override
@@ -89,6 +120,7 @@ public ResumableUploadFuture<ResponseT> futureCall(
89120
return ResumableUploadFutureImpl.create(
90121
startFuture,
91122
client.uploadChunkCallable(),
123+
retryingQueryCallable,
92124
payload,
93125
effectiveSettings,
94126
clientContext.getDefaultCallContext(),

0 commit comments

Comments
 (0)