feat(gax): add recoverable error query and buffer realignment loop

Introduces a query-and-realign loop when encountering recoverable protocol errors during chunk uploads. 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,
@@ -86,6 +90,9 @@ public ResumableUploadCallableImpl(
     this.retryingUploadChunkCallable =
         createRetryingCallable(
             client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD, clientContext);
+    this.retryingQueryCallable =
+        createRetryingCallable(
+            client.queryStatusCallable(), ResumableUploadCommand.QUERY, clientContext);
   }
 
   @Override
@@ -109,6 +116,7 @@ public ResumableUploadFuture<ResponseT> futureCall(
     return ResumableUploadFutureImpl.create(
         startFuture,
         retryingUploadChunkCallable,
+        retryingQueryCallable,
         payload,
         effectiveSettings,
         clientContext.getDefaultCallContext());

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;
@@ -50,6 +53,9 @@
 /**
  * Coordinates chunk transmission steps of a resumable upload session.
  *
+ * <p>Expects {@code uploadChunkCallable} and {@code queryStatusCallable} to be pre-wrapped in
+ * retrying callables that handle transient errors.
+ *
  * @param <ResponseT> the type of the final response message returned once the upload completes
  */
 @InternalApi
@@ -62,6 +68,8 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
 
   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;
@@ -70,12 +78,15 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
 
   ResumableUploadChunkCoordinator(
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       String uploadUrl,
       InputStream payload,
       int chunkSize,
       ApiCallContext callContext) {
     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");
     this.callContext = checkNotNull(callContext, "callContext must not be null");
@@ -96,60 +107,85 @@ 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;
+      dispatchCurrentChunk();
+    } catch (Throwable t) {
+      result.setException(t);
     }
+  }
 
-    // Determine if this is the final chunk and build the chunk request.
-    ChunkUploadRequest chunkRequest =
-        ChunkUploadRequest.newBuilder()
-            .setUploadUrl(uploadUrl)
-            .setPayload(buffer.getBuffer())
-            .setPayloadLength(buffer.getPayloadLength())
-            .setOffset(buffer.getBufferBaseOffset())
-            .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();
+  private void dispatchCurrentChunk() {
+    if (result.isDone()) {
+      return;
+    }
     try {
+      ChunkUploadRequest chunkRequest = buildCurrentChunkRequest();
       ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
           uploadChunkCallable.futureCall(chunkRequest, callContext);
       this.currentChunkFuture = chunkFuture;
       if (result.isCancelled()) {
         chunkFuture.cancel(true);
         return;
       }
-
       ApiFutures.addCallback(
           chunkFuture,
           new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
             @Override
             public void onSuccess(ChunkUploadResponse<ResponseT> response) {
-              if (result.isDone()) {
+              if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+                recover();
+              } else {
+                onChunkUploaded(response);
+              }
+            }
+
+            @Override
+            public void onFailure(Throwable t) {
+              if (t instanceof CancellationException || result.isDone()) {
                 return;
               }
-              long nextOffset = currentOffset + chunkLength;
-              if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
-                result.set(response.getResponse());
-              } else if (isFinal) {
-                result.setException(
-                    new IllegalStateException(
-                        "Upload stream ended and final chunk was transmitted, but server returned"
-                            + " incomplete status for upload URL: "
-                            + uploadUrl));
+              Category category =
+                  ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD);
+              if (category == Category.RECOVERABLE) {
+                recover();
               } else {
-                chunkExecutor.execute(() -> transmitChunk(nextOffset));
+                // Category.TRANSIENT errors reaching here have already exhausted their retry budget
+                // in the underlying RetryingCallable and become fatal per protocol specification.
+                result.setException(t);
+              }
+            }
+          },
+          chunkExecutor);
+    } catch (Throwable t) {
+      result.setException(t);
+    }
+  }
+
+  private void recover() {
+    if (result.isDone()) {
+      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) {
+              try {
+                handleQueryResponse(queryResponse);
+              } catch (Throwable t) {
+                result.setException(t);
               }
             }
 
@@ -161,10 +197,78 @@ public void onFailure(Throwable t) {
               result.setException(t);
             }
           },
-          MoreExecutors.directExecutor());
+          chunkExecutor);
     } catch (Throwable t) {
       result.setException(t);
     }
   }
-}
 
+  private void handleQueryResponse(QueryStatusResponse<ResponseT> queryResponse) throws IOException {
+    if (result.isDone()) {
+      return;
+    }
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+      throw protocolViolation(
+          "Query status response missing X-Goog-Upload-Status header for upload URL: " + uploadUrl);
+    }
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) {
+      onChunkUploaded(
+          ChunkUploadResponse.create(ResumableUploadStatus.FINAL, queryResponse.getResponse()));
+      return;
+    }
+    Long committedOffset = queryResponse.getCommittedOffset();
+    if (committedOffset == null) {
+      throw protocolViolation(
+          "Incomplete query status response did not include a committed offset for upload URL: "
+              + uploadUrl);
+    }
+    buffer.realignTo(committedOffset);
+    dispatchCurrentChunk();
+  }
+
+  private void onChunkUploaded(ChunkUploadResponse<ResponseT> response) {
+    if (result.isDone()) {
+      return;
+    }
+    long nextOffset = buffer.getBufferBaseOffset() + buffer.getPayloadLength();
+    if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
+      result.set(response.getResponse());
+    } 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 {
+      chunkExecutor.execute(() -> transmitChunk(nextOffset));
+    }
+  }
+
+  private ChunkUploadRequest buildCurrentChunkRequest() {
+    // Determine if this is the final chunk and build the chunk request.
+    return ChunkUploadRequest.newBuilder()
+        .setUploadUrl(uploadUrl)
+        .setPayload(buffer.getBuffer())
+        .setPayloadLength(buffer.getPayloadLength())
+        .setOffset(buffer.getBufferBaseOffset())
+        .setFinal(buffer.isFinal())
+        .build();
+  }
+
+  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 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 ApiCallContext callContext;
@@ -85,12 +89,18 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   static <ResponseT> ResumableUploadFutureImpl<ResponseT> create(
       ApiFuture<ResumableUploadSession> startFuture,
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       InputStream payload,
       ResumableUploadCallSettings settings,
       ApiCallContext callContext) {
     ResumableUploadFutureImpl<ResponseT> future =
         new ResumableUploadFutureImpl<>(
-            startFuture, uploadChunkCallable, payload, settings, callContext);
+            startFuture,
+            uploadChunkCallable,
+            queryStatusCallable,
+            payload,
+            settings,
+            callContext);
     try {
       future.start();
     } catch (Throwable t) {
@@ -102,12 +112,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,
       ApiCallContext callContext) {
     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 +141,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();