Skip to content

Commit 8456db3

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 c6c28b1 commit 8456db3

11 files changed

Lines changed: 1265 additions & 116 deletions

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

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

‎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)