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.

Author: whowes

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

@@ -0,0 +1,306 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.rpc;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.api.core.ApiFuture;
+import com.google.api.core.ApiFutureCallback;
+import com.google.api.core.ApiFutures;
+import com.google.api.core.SettableApiFuture;
+import com.google.api.gax.resumable.ChunkUploadRequest;
+import com.google.api.gax.resumable.ChunkUploadResponse;
+import com.google.api.gax.resumable.QueryStatusRequest;
+import com.google.api.gax.resumable.QueryStatusResponse;
+import com.google.api.gax.resumable.ResumableUploadStatus;
+import com.google.api.gax.retrying.RetryingFuture;
+import com.google.common.util.concurrent.MoreExecutors;
+import java.time.Duration;
+import java.util.concurrent.Callable;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * A {@link Callable} representing an attempt to transmit a single chunk in a resumable upload
+ * session. Used with {@link com.google.api.gax.retrying.ScheduledRetryingExecutor}.
+ *
+ * <p>Execution follows the standard attempt template with pre-attempt recovery handling. When the
+ * previous attempt failed with a Category 2 (recoverable) error or missing status header, {@code
+ * prepareAttempt} queries session status, realigns the buffer window, tops up from the stream, and
+ * dispatches the chunk upload request. The callable never blocks on {@code .get()}; results and
+ * cancellations propagate asynchronously.
+ *
+ * @param <ResponseT> the type of the final response message once the upload completes
+ */
+@NullMarked
+class ChunkAttemptCallable<ResponseT> implements Callable<ChunkUploadResponse<ResponseT>> {
+
+  private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
+      uploadChunkCallable;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      queryStatusCallable;
+  private final RewindableStreamBuffer buffer;
+  private final String uploadUrl;
+  private final ApiCallContext originalCallContext;
+
+  private volatile ChunkUploadRequest currentRequest;
+  private volatile ResumableUploadCommand currentCommand;
+
+  private volatile @Nullable RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture;
+  private volatile @Nullable ApiFuture<?> inFlightFuture;
+  private volatile @Nullable Throwable lastFailure;
+  private volatile @Nullable ChunkUploadResponse<ResponseT> lastResponse;
+
+  ChunkAttemptCallable(
+      UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
+      RewindableStreamBuffer buffer,
+      String uploadUrl,
+      ChunkUploadRequest request,
+      ApiCallContext callContext,
+      ResumableUploadCommand command) {
+    this.uploadChunkCallable =
+        checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
+    this.queryStatusCallable =
+        checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
+    this.buffer = checkNotNull(buffer, "buffer must not be null");
+    this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
+    this.currentRequest = checkNotNull(request, "request must not be null");
+    this.originalCallContext = checkNotNull(callContext, "callContext must not be null");
+    this.currentCommand = checkNotNull(command, "command must not be null");
+  }
+
+  void setRetryingFuture(RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture) {
+    this.retryingFuture = checkNotNull(retryingFuture, "retryingFuture must not be null");
+  }
+
+  ResumableUploadCommand getCurrentCommand() {
+    return currentCommand;
+  }
+
+  private boolean needsRecovery() {
+    if (lastFailure != null) {
+      ResumableUploadErrorClassifier.Category category =
+          ResumableUploadErrorClassifier.classify(lastFailure, currentCommand);
+      return category == ResumableUploadErrorClassifier.Category.RECOVERABLE;
+    }
+    if (lastResponse != null && lastResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+      ResumableUploadErrorClassifier.Category category =
+          ResumableUploadErrorClassifier.classifyMissingStatusHeader(currentCommand);
+      return category == ResumableUploadErrorClassifier.Category.RECOVERABLE;
+    }
+    return false;
+  }
+
+  private void failAttempt(
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture, Throwable t) {
+    lastFailure = t;
+    lastResponse = null;
+    attemptFuture.setException(t);
+  }
+
+  /**
+   * Pre-attempt recovery step invoked before transmitting an attempt when the previous attempt
+   * encountered a Category 2 (recoverable) error or missing status header.
+   */
+  private void prepareAttempt(
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
+      ApiCallContext attemptContext,
+      RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    this.currentCommand = ResumableUploadCommand.QUERY;
+    QueryStatusRequest queryRequest = QueryStatusRequest.create(uploadUrl);
+    ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
+        queryStatusCallable.futureCall(queryRequest, attemptContext);
+    if (queryFuture == null) {
+      failAttempt(
+          attemptFuture, new IllegalStateException("queryStatusCallable returned a null future"));
+      return;
+    }
+    this.inFlightFuture = queryFuture;
+
+    ApiFutures.addCallback(
+        queryFuture,
+        new ApiFutureCallback<QueryStatusResponse<ResponseT>>() {
+          @Override
+          public void onSuccess(QueryStatusResponse<ResponseT> queryResponse) {
+            handleQuerySuccess(queryResponse, attemptFuture, attemptContext, currentRetryingFuture);
+          }
+
+          @Override
+          public void onFailure(Throwable t) {
+            failAttempt(attemptFuture, t);
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  private void handleQuerySuccess(
+      QueryStatusResponse<ResponseT> queryResponse,
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
+      ApiCallContext attemptContext,
+      RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    if (currentRetryingFuture.isDone()) {
+      return;
+    }
+
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+      failAttempt(
+          attemptFuture,
+          UploadErrors.protocolViolation(
+              "Query status response missing X-Goog-Upload-Status header for upload URL: "
+                  + uploadUrl));
+      return;
+    }
+
+    // Server already finalized the upload.
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) {
+      ChunkUploadResponse<ResponseT> response =
+          ChunkUploadResponse.create(ResumableUploadStatus.FINAL, queryResponse.getResponse());
+      lastFailure = null;
+      lastResponse = response;
+      attemptFuture.set(response);
+      return;
+    }
+
+    // Incomplete query response with null committed offset violates the protocol invariant.
+    Long committedOffset = queryResponse.getCommittedOffset();
+    if (committedOffset == null) {
+      failAttempt(
+          attemptFuture,
+          UploadErrors.protocolViolation(
+              "Incomplete query status response did not include a committed offset for upload URL: "
+                  + uploadUrl));
+      return;
+    }
+
+    // Normal path: realign buffer to committedOffset, compact and top up.
+    try {
+      buffer.realignTo(committedOffset);
+    } catch (Throwable e) {
+      failAttempt(attemptFuture, e);
+      return;
+    }
+
+    // Determine the upload command for the realigned buffer.
+    // Preserve upload,finalize for a trailing partial after realignment.
+    ResumableUploadCommand realignedCommand;
+    if (buffer.isFinal()) {
+      realignedCommand =
+          buffer.isEmpty()
+              ? ResumableUploadCommand.FINALIZE
+              : ResumableUploadCommand.UPLOAD_FINALIZE;
+    } else {
+      realignedCommand = ResumableUploadCommand.UPLOAD;
+    }
+
+    ChunkUploadRequest realignedRequest =
+        ChunkUploadRequest.newBuilder()
+            .setUploadUrl(uploadUrl)
+            .setPayload(buffer.getBuffer())
+            .setPayloadLength(buffer.getPayloadLength())
+            .setOffset(buffer.getBufferBaseOffset())
+            .setFinal(buffer.isFinal())
+            .build();
+
+    this.currentRequest = realignedRequest;
+    this.currentCommand = realignedCommand;
+
+    dispatchChunkUpload(attemptFuture, attemptContext, currentRetryingFuture);
+  }
+
+  private void dispatchChunkUpload(
+      SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture,
+      ApiCallContext attemptContext,
+      RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture) {
+    ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
+        uploadChunkCallable.futureCall(currentRequest, attemptContext);
+    this.inFlightFuture = chunkFuture;
+
+    ApiFutures.addCallback(
+        chunkFuture,
+        new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
+          @Override
+          public void onSuccess(ChunkUploadResponse<ResponseT> response) {
+            lastFailure = null;
+            lastResponse = response;
+            attemptFuture.set(response);
+          }
+
+          @Override
+          public void onFailure(Throwable t) {
+            failAttempt(attemptFuture, t);
+          }
+        },
+        MoreExecutors.directExecutor());
+  }
+
+  @Override
+  public @Nullable ChunkUploadResponse<ResponseT> call() {
+    RetryingFuture<ChunkUploadResponse<ResponseT>> currentRetryingFuture =
+        checkNotNull(retryingFuture, "retryingFuture must be set before call()");
+    ApiCallContext attemptContext = originalCallContext;
+
+    Duration rpcTimeout = currentRetryingFuture.getAttemptSettings().getRpcTimeoutDuration();
+    if (!rpcTimeout.isZero() && attemptContext.getTimeoutDuration() == null) {
+      attemptContext = attemptContext.withTimeoutDuration(rpcTimeout);
+    }
+
+    SettableApiFuture<ChunkUploadResponse<ResponseT>> attemptFuture = SettableApiFuture.create();
+    currentRetryingFuture.setAttemptFuture(attemptFuture);
+
+    if (currentRetryingFuture.isDone()) {
+      return null;
+    }
+
+    currentRetryingFuture.addListener(
+        () -> {
+          if (currentRetryingFuture.isCancelled()) {
+            ApiFuture<?> inFlight = inFlightFuture;
+            if (inFlight != null) {
+              inFlight.cancel(true);
+            }
+            attemptFuture.cancel(true);
+          }
+        },
+        MoreExecutors.directExecutor());
+
+    try {
+      if (needsRecovery()) {
+        prepareAttempt(attemptFuture, attemptContext, currentRetryingFuture);
+      } else {
+        dispatchChunkUpload(attemptFuture, attemptContext, currentRetryingFuture);
+      }
+    } catch (Throwable t) {
+      failAttempt(attemptFuture, t);
+    }
+
+    return null;
+  }
+}

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

@@ -35,9 +35,16 @@
 import com.google.api.core.ApiFutures;
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
+import com.google.api.gax.resumable.QueryStatusRequest;
+import com.google.api.gax.resumable.QueryStatusResponse;
 import com.google.api.gax.resumable.ResumableUploadClient;
 import com.google.api.gax.resumable.ResumableUploadSession;
+import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
+import com.google.api.gax.retrying.RetryAlgorithm;
+import com.google.api.gax.retrying.RetrySettings;
+import com.google.api.gax.retrying.ScheduledRetryingExecutor;
 import java.io.InputStream;
+import java.time.Duration;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -54,9 +61,22 @@
 public class ResumableUploadCallableImpl<RequestT, ResponseT>
     extends ResumableUploadCallable<RequestT, ResponseT> {
 
+  static final RetrySettings DEFAULT_QUERY_RETRY_SETTINGS =
+      RetrySettings.newBuilder()
+          .setInitialRetryDelayDuration(Duration.ofMillis(100))
+          .setRetryDelayMultiplier(1.3)
+          .setMaxRetryDelayDuration(Duration.ofMinutes(1))
+          .setInitialRpcTimeoutDuration(Duration.ofSeconds(30))
+          .setRpcTimeoutMultiplier(1.0)
+          .setMaxRpcTimeoutDuration(Duration.ofSeconds(30))
+          .setTotalTimeoutDuration(Duration.ofMinutes(5))
+          .build();
+
   private final ResumableUploadClient<RequestT, ResponseT> client;
   private final ResumableUploadCallSettings defaultCallSettings;
   private final ClientContext clientContext;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      retryingQueryCallable;
 
   public ResumableUploadCallableImpl(
       ResumableUploadClient<RequestT, ResponseT> client,
@@ -66,6 +86,17 @@ public ResumableUploadCallableImpl(
     this.defaultCallSettings =
         checkNotNull(defaultCallSettings, "defaultCallSettings must not be null");
     this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
+
+    RetryAlgorithm<QueryStatusResponse<ResponseT>> queryRetryAlgorithm =
+        new RetryAlgorithm<>(
+            new ResumableUploadResultRetryAlgorithm<>(ResumableUploadCommand.QUERY),
+            new ExponentialRetryAlgorithm(DEFAULT_QUERY_RETRY_SETTINGS, clientContext.getClock()));
+
+    this.retryingQueryCallable =
+        new RetryingCallable<>(
+            clientContext.getDefaultCallContext(),
+            checkNotNull(client.queryStatusCallable(), "queryStatusCallable must not be null"),
+            new ScheduledRetryingExecutor<>(queryRetryAlgorithm, clientContext.getExecutor()));
   }
 
   @Override
@@ -89,6 +120,7 @@ public ResumableUploadFuture<ResponseT> futureCall(
     return ResumableUploadFutureImpl.create(
         startFuture,
         client.uploadChunkCallable(),
+        retryingQueryCallable,
         payload,
         effectiveSettings,
         clientContext.getDefaultCallContext(),