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/ResumableUploadCallableImpl.java

@@ -37,6 +37,8 @@
 import com.google.api.core.InternalApi;
 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.ResumableUploadClient;
 import com.google.api.gax.resumable.ResumableUploadSession;
 import com.google.api.gax.retrying.ExponentialRetryAlgorithm;
@@ -74,6 +76,8 @@ public class ResumableUploadCallableImpl<RequestT, ResponseT>
   private final ClientContext clientContext;
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       retryingUploadChunkCallable;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      retryingQueryCallable;
 
   public ResumableUploadCallableImpl(
       ResumableUploadClient<RequestT, ResponseT> client,
@@ -89,6 +93,8 @@ public ResumableUploadCallableImpl(
             .build();
     this.retryingUploadChunkCallable =
         createRetryingCallable(client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD);
+    this.retryingQueryCallable =
+        createRetryingCallable(client.queryStatusCallable(), ResumableUploadCommand.QUERY);
   }
 
   @Override
@@ -110,7 +116,12 @@ public ResumableUploadFuture<ResponseT> futureCall(
     }
 
     return ResumableUploadFutureImpl.create(
-        startFuture, retryingUploadChunkCallable, payload, effectiveSettings, clientContext);
+        startFuture,
+        retryingUploadChunkCallable,
+        retryingQueryCallable,
+        payload,
+        effectiveSettings,
+        clientContext);
   }
 
   @Override

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

@@ -38,7 +38,10 @@
 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.rpc.ResumableUploadErrorClassifier.Category;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.io.IOException;
 import java.io.InputStream;
@@ -56,25 +59,44 @@
 @NullMarked
 final class ResumableUploadChunkCoordinator<ResponseT> {
 
+  private static final StatusCode FAILED_PRECONDITION_STATUS_CODE =
+      new StatusCode() {
+        @Override
+        public StatusCode.Code getCode() {
+          return StatusCode.Code.FAILED_PRECONDITION;
+        }
+
+        @Override
+        public @Nullable Object getTransportCode() {
+          return null;
+        }
+      };
+
   private final Executor chunkExecutor =
       MoreExecutors.newSequentialExecutor(MoreExecutors.directExecutor());
 
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       uploadChunkCallable;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      queryStatusCallable;
   private final String uploadUrl;
   private final RewindableStreamBuffer buffer;
   private final ApiCallContext callContext;
   private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
   private volatile @Nullable ApiFuture<?> currentChunkFuture;
+  private int recoveryAttempts;
 
   ResumableUploadChunkCoordinator(
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       String uploadUrl,
       InputStream payload,
       int chunkSize,
       ClientContext clientContext) {
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
+    this.queryStatusCallable =
+        checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
     this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
     checkNotNull(payload, "payload must not be null");
     checkNotNull(clientContext, "clientContext must not be null");
@@ -96,20 +118,26 @@ ApiFuture<ResponseT> start() {
   }
 
   private void transmitChunk(long currentOffset) {
-    // Abort if the session was already completed or canceled.
     if (result.isDone()) {
       return;
     }
 
-    // Read the next chunk slice from the payload stream.
     try {
       buffer.fill(currentOffset);
     } catch (IOException e) {
       result.setException(e);
       return;
     }
 
-    // Determine if this is the final chunk and build the chunk request.
+    recoveryAttempts = 0;
+    sendCurrentBuffer();
+  }
+
+  private void sendCurrentBuffer() {
+    if (result.isDone()) {
+      return;
+    }
+
     ChunkUploadRequest chunkRequest =
         ChunkUploadRequest.newBuilder()
             .setUploadUrl(uploadUrl)
@@ -119,9 +147,6 @@ private void transmitChunk(long currentOffset) {
             .setFinal(buffer.isFinal())
             .build();
 
-    // Dispatch the chunk upload call and register the in-flight future for cancellation.
-    long chunkLength = chunkRequest.getPayloadLength();
-    boolean isFinal = chunkRequest.isFinal();
     try {
       ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
           uploadChunkCallable.futureCall(chunkRequest, callContext);
@@ -139,16 +164,18 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
               if (result.isDone()) {
                 return;
               }
-              long nextOffset = currentOffset + chunkLength;
-              if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
+              if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+                chunkExecutor.execute(() -> recoverChunk(null));
+              } else if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
                 result.set(response.getResponse());
-              } else if (isFinal) {
+              } else if (buffer.isFinal()) {
                 result.setException(
                     new IllegalStateException(
                         "Upload stream ended and final chunk was transmitted, but server returned"
                             + " incomplete status for upload URL: "
                             + uploadUrl));
               } else {
+                long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
                 chunkExecutor.execute(() -> transmitChunk(nextOffset));
               }
             }
@@ -158,6 +185,11 @@ public void onFailure(Throwable t) {
               if (t instanceof CancellationException || result.isDone()) {
                 return;
               }
+              if (ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD)
+                  == Category.RECOVERABLE) {
+                chunkExecutor.execute(() -> recoverChunk(t));
+                return;
+              }
               result.setException(t);
             }
           },
@@ -166,4 +198,89 @@ public void onFailure(Throwable t) {
       result.setException(t);
     }
   }
+
+  private void recoverChunk(@Nullable Throwable failure) {
+    if (result.isDone()) {
+      return;
+    }
+    int maxAttempts =
+        callContext.getRetrySettings() != null
+            ? callContext.getRetrySettings().getMaxAttempts()
+            : 5;
+    if (++recoveryAttempts >= maxAttempts) {
+      result.setException(
+          failure != null
+              ? failure
+              : protocolViolation(
+                  "Chunk upload response missing X-Goog-Upload-Status header for upload URL: "
+                      + uploadUrl));
+      return;
+    }
+
+    try {
+      ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
+          queryStatusCallable.futureCall(QueryStatusRequest.create(uploadUrl), callContext);
+      this.currentChunkFuture = queryFuture;
+      if (result.isCancelled()) {
+        queryFuture.cancel(true);
+        return;
+      }
+
+      ApiFutures.addCallback(
+          queryFuture,
+          new ApiFutureCallback<QueryStatusResponse<ResponseT>>() {
+            @Override
+            public void onSuccess(QueryStatusResponse<ResponseT> queryResponse) {
+              chunkExecutor.execute(() -> handleQueryResponse(queryResponse));
+            }
+
+            @Override
+            public void onFailure(Throwable t) {
+              if (t instanceof CancellationException || result.isDone()) {
+                return;
+              }
+              result.setException(t);
+            }
+          },
+          MoreExecutors.directExecutor());
+    } catch (Throwable t) {
+      result.setException(t);
+    }
+  }
+
+  private void handleQueryResponse(QueryStatusResponse<ResponseT> queryResponse) {
+    if (result.isDone()) {
+      return;
+    }
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+      result.setException(
+          protocolViolation(
+              "Query status response missing X-Goog-Upload-Status header for upload URL: "
+                  + uploadUrl));
+      return;
+    }
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) {
+      result.set(queryResponse.getResponse());
+      return;
+    }
+    Long committedOffset = queryResponse.getCommittedOffset();
+    if (committedOffset == null) {
+      result.setException(
+          protocolViolation(
+              "Incomplete query status response did not include a committed offset for upload URL: "
+                  + uploadUrl));
+      return;
+    }
+    try {
+      buffer.realignTo(committedOffset);
+    } catch (Throwable t) {
+      result.setException(t);
+      return;
+    }
+    sendCurrentBuffer();
+  }
+
+  private static FailedPreconditionException protocolViolation(String message) {
+    return new FailedPreconditionException(message, null, FAILED_PRECONDITION_STATUS_CODE, false);
+  }
 }

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

@@ -38,6 +38,8 @@
 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.ResumableUploadSession;
 import com.google.common.util.concurrent.MoreExecutors;
 import com.google.errorprone.annotations.concurrent.GuardedBy;
@@ -65,6 +67,8 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   private final ApiFuture<ResumableUploadSession> startFuture;
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       uploadChunkCallable;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      queryStatusCallable;
   private final InputStream payload;
   private final ResumableUploadCallSettings settings;
   private final ClientContext clientContext;
@@ -75,22 +79,21 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   @GuardedBy("lock")
   private @Nullable ApiFuture<?> inFlightFuture;
 
-  /**
-   * Creates and initiates a new resumable upload future tracking session initiation and chunk
-   * streaming.
-   *
-   * <p>The provided {@code payload} stream is managed by the returned future and will be closed
-   * automatically upon completion, failure, or cancellation.
-   */
   static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
       ApiFuture<ResumableUploadSession> startFuture,
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
       ClientContext clientContext) {
     ResumableUploadFutureImpl<ResponseT> future =
         new ResumableUploadFutureImpl<>(
-            startFuture, uploadChunkCallable, payload, settings, clientContext);
+            startFuture,
+            uploadChunkCallable,
+            queryStatusCallable,
+            payload,
+            settings,
+            clientContext);
     try {
       future.start();
     } catch (Throwable t) {
@@ -102,12 +105,15 @@ static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
   private ResumableUploadFutureImpl(
       ApiFuture<ResumableUploadSession> startFuture,
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
       ClientContext clientContext) {
     this.startFuture = checkNotNull(startFuture, "startFuture must not be null");
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
+    this.queryStatusCallable =
+        checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
     this.payload = checkNotNull(payload, "payload must not be null");
     this.settings = checkNotNull(settings, "settings must not be null");
     checkArgument(settings.getChunkSize() > 0, "chunkSize must be > 0");
@@ -128,6 +134,7 @@ public void onSuccess(ResumableUploadSession session) {
             ResumableUploadChunkCoordinator<ResponseT> coordinator =
                 new ResumableUploadChunkCoordinator<>(
                     uploadChunkCallable,
+                    queryStatusCallable,
                     uploadSessionUrl,
                     payload,
                     settings.getChunkSize(),

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/CallableTest.java

@@ -216,6 +216,8 @@ void testResumableUploadCallable() {
         mock(ResumableUploadClient.class, Mockito.withSettings().withoutAnnotations());
     when(uploadClient.uploadChunkCallable())
         .thenReturn(mock(UnaryCallable.class, Mockito.withSettings().withoutAnnotations()));
+    when(uploadClient.queryStatusCallable())
+        .thenReturn(mock(UnaryCallable.class, Mockito.withSettings().withoutAnnotations()));
     ResumableUploadCallSettings settings =
         ResumableUploadCallSettings.newBuilder().setChunkSize(1024).build();