4141import com .google .api .gax .resumable .QueryStatusRequest ;
4242import com .google .api .gax .resumable .QueryStatusResponse ;
4343import com .google .api .gax .resumable .ResumableUploadStatus ;
44+ import com .google .api .gax .retrying .RetrySettings ;
4445import com .google .api .gax .rpc .ResumableUploadErrorClassifier .Category ;
4546import com .google .common .util .concurrent .MoreExecutors ;
4647import java .io .IOException ;
4748import java .io .InputStream ;
49+ import java .time .Duration ;
4850import java .util .concurrent .CancellationException ;
4951import java .util .concurrent .Executor ;
5052import org .jspecify .annotations .NullMarked ;
@@ -72,7 +74,10 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
7274 queryStatusCallable ;
7375 private final String uploadUrl ;
7476 private final RewindableStreamBuffer buffer ;
77+ private final ResumableUploadCallSettings settings ;
78+ private final long deadlineNanos ;
7579 private final ApiCallContext callContext ;
80+ private final ClientContext clientContext ;
7681 private final SettableApiFuture <ResponseT > result = SettableApiFuture .create ();
7782 private volatile @ Nullable ApiFuture <?> currentChunkFuture ;
7883
@@ -81,16 +86,20 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
8186 UnaryCallable <QueryStatusRequest , QueryStatusResponse <ResponseT >> queryStatusCallable ,
8287 String uploadUrl ,
8388 InputStream payload ,
84- int chunkSize ,
85- ApiCallContext callContext ) {
89+ ResumableUploadCallSettings settings ,
90+ long deadlineNanos ,
91+ ClientContext clientContext ) {
8692 this .uploadChunkCallable =
8793 checkNotNull (uploadChunkCallable , "uploadChunkCallable must not be null" );
8894 this .queryStatusCallable =
8995 checkNotNull (queryStatusCallable , "queryStatusCallable must not be null" );
9096 this .uploadUrl = checkNotNull (uploadUrl , "uploadUrl must not be null" );
9197 checkNotNull (payload , "payload must not be null" );
92- this .callContext = checkNotNull (callContext , "callContext must not be null" );
93- this .buffer = new RewindableStreamBuffer (payload , chunkSize , uploadUrl );
98+ this .settings = checkNotNull (settings , "settings must not be null" );
99+ this .deadlineNanos = deadlineNanos ;
100+ this .clientContext = checkNotNull (clientContext , "clientContext must not be null" );
101+ this .callContext = clientContext .getDefaultCallContext ();
102+ this .buffer = new RewindableStreamBuffer (payload , settings .getChunkSize (), uploadUrl );
94103 }
95104
96105 ApiFuture <ResponseT > start () {
@@ -126,7 +135,7 @@ private void dispatchCurrentChunk() {
126135 ChunkUploadRequest chunkRequest = buildCurrentChunkRequest ();
127136 // Dispatch the chunk upload call and register the in-flight future for cancellation.
128137 ApiFuture <ChunkUploadResponse <ResponseT >> chunkFuture =
129- uploadChunkCallable .futureCall (chunkRequest , callContext );
138+ uploadChunkCallable .futureCall (chunkRequest , chunkCallContext () );
130139 if (!trackInFlight (chunkFuture )) {
131140 return ;
132141 }
@@ -171,7 +180,7 @@ private void recover() {
171180 try {
172181 // Dispatch the query status call and register the in-flight future for cancellation.
173182 ApiFuture <QueryStatusResponse <ResponseT >> queryFuture =
174- queryStatusCallable .futureCall (QueryStatusRequest .create (uploadUrl ), callContext );
183+ queryStatusCallable .futureCall (QueryStatusRequest .create (uploadUrl ), queryCallContext () );
175184 if (!trackInFlight (queryFuture )) {
176185 return ;
177186 }
@@ -273,6 +282,60 @@ public StatusCode.Code getCode() {
273282 }
274283 };
275284
285+ /**
286+ * Returns the call context for a chunk upload, bounding its retry cycle and each of its attempts
287+ * to the chunk's local deadline: at most half the global timeout, never past the global deadline.
288+ */
289+ private ApiCallContext chunkCallContext () {
290+ Duration remaining = remainingGlobalTimeout ();
291+ if (remaining == null ) {
292+ return callContext ;
293+ }
294+ Duration deadline = min (settings .getGlobalTimeout ().dividedBy (2 ), remaining );
295+ return boundedCallContext (
296+ deadline ,
297+ callContext .getRetrySettings ().toBuilder ()
298+ .setInitialRpcTimeoutDuration (deadline )
299+ .setRpcTimeoutMultiplier (1.0 )
300+ .setMaxRpcTimeoutDuration (deadline ));
301+ }
302+
303+ /** Returns the call context for a status query, bounded by the remaining global timeout. */
304+ private ApiCallContext queryCallContext () {
305+ Duration remaining = remainingGlobalTimeout ();
306+ if (remaining == null ) {
307+ return callContext ;
308+ }
309+ RetrySettings baseSettings = callContext .getRetrySettings ();
310+ return boundedCallContext (
311+ remaining ,
312+ baseSettings .toBuilder ()
313+ .setInitialRpcTimeoutDuration (
314+ min (baseSettings .getInitialRpcTimeoutDuration (), remaining ))
315+ .setMaxRpcTimeoutDuration (min (baseSettings .getMaxRpcTimeoutDuration (), remaining )));
316+ }
317+
318+ /**
319+ * Returns the time left before the global upload deadline, or null if deadlines cannot be derived
320+ * because no global timeout or no base retry settings were configured.
321+ */
322+ private @ Nullable Duration remainingGlobalTimeout () {
323+ if (settings .getGlobalTimeout () == null || callContext .getRetrySettings () == null ) {
324+ return null ;
325+ }
326+ return Duration .ofNanos (Math .max (1L , deadlineNanos - clientContext .getClock ().nanoTime ()));
327+ }
328+
329+ private ApiCallContext boundedCallContext (Duration deadline , RetrySettings .Builder settings ) {
330+ return callContext
331+ .withRetrySettings (settings .setTotalTimeoutDuration (deadline ).build ())
332+ .withTimeoutDuration (deadline );
333+ }
334+
335+ private static Duration min (Duration left , Duration right ) {
336+ return left .compareTo (right ) <= 0 ? left : right ;
337+ }
338+
276339 private static FailedPreconditionException protocolViolation (String message ) {
277340 return new FailedPreconditionException (message , null , FAILED_PRECONDITION_STATUS_CODE , false );
278341 }
0 commit comments