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

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;
@@ -50,21 +53,26 @@
 /**
  * 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
 @NullMarked
 final class ResumableUploadChunkCoordinator<ResponseT> {
 
   /**
-   * Serializes every read and write of {@link #buffer}, so that chunk transmission and
-   * its asynchronous continuations never touch the buffer window concurrently.
+   * Serializes every read and write of {@link #buffer}, so that chunk transmission, recovery, and
+   * their asynchronous continuations never touch the buffer window concurrently.
    */
   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;
@@ -73,12 +81,15 @@
 
   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");
@@ -100,66 +111,39 @@
   }
 
   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 {
-      ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
-          uploadChunkCallable.futureCall(chunkRequest, callContext);
+      ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture = executeChunkWithRecovery();
       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()) {
-                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));
-              } else {
-                chunkExecutor.execute(() -> transmitChunk(nextOffset));
-              }
+              chunkExecutor.execute(() -> onChunkUploaded(response));
             }
 
             @Override
             public void onFailure(Throwable t) {
-              if (t instanceof CancellationException || result.isDone()) {
+              if (t instanceof CancellationException) {
                 return;
               }
               result.setException(t);
@@ -170,4 +154,108 @@
       result.setException(t);
     }
   }
+
+  private ApiFuture<ChunkUploadResponse<ResponseT>> executeChunkWithRecovery() {
+    ApiFuture<ChunkUploadResponse<ResponseT>> transmitted =
+        ApiFutures.catchingAsync(
+            uploadChunkCallable.futureCall(buildCurrentChunkRequest(), callContext),
+            Throwable.class,
+            t -> {
+              Category category =
+                  ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD);
+              if (category == Category.RECOVERABLE) {
+                return recover();
+              }
+              // Category.TRANSIENT errors reaching here have already exhausted their retry budget
+              // in the underlying RetryingCallable and become fatal per protocol specification.
+              return ApiFutures.immediateFailedFuture(t);
+            },
+            chunkExecutor);
+    return ApiFutures.transformAsync(
+        transmitted,
+        response -> {
+          if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+            return recover();
+          }
+          return ApiFutures.immediateFuture(response);
+        },
+        chunkExecutor);
+  }
+
+  private ApiFuture<ChunkUploadResponse<ResponseT>> recover() {
+    return ApiFutures.transformAsync(
+        queryStatusCallable.futureCall(QueryStatusRequest.create(uploadUrl), callContext),
+        this::resumeFrom,
+        chunkExecutor);
+  }
+
+  private ApiFuture<ChunkUploadResponse<ResponseT>> resumeFrom(
+      QueryStatusResponse<ResponseT> queryResponse) throws IOException {
+    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) {
+      return ApiFutures.immediateFuture(
+          ChunkUploadResponse.create(ResumableUploadStatus.FINAL, queryResponse.getResponse()));
+    }
+    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);
+    return executeChunkWithRecovery();
+  }
+
+  private void onChunkUploaded(ChunkUploadResponse<ResponseT> response) {
+    if (result.isDone()) {
+      return;
+    }
+    try {
+      if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
+        result.set(response.getResponse());
+      } else if (buffer.isFinal()) {
+        throw 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();
+        transmitChunk(nextOffset);
+      }
+    } catch (Throwable t) {
+      result.setException(t);
+    }
+  }
+
+  private ChunkUploadRequest buildCurrentChunkRequest() {
+    return ChunkUploadRequest.newBuilder()
+        .setUploadUrl(uploadUrl)
+        .setPayload(buffer.getBuffer())
+        .setPayloadLength(buffer.getPayloadLength())
+        .setOffset(buffer.getBufferBaseOffset())
+        .setFinal(buffer.isFinal())
+        .build();
+  }
+
+  private static FailedPreconditionException protocolViolation(String message) {
+    return new FailedPreconditionException(
+        message,
+        null,
+        new StatusCode() {
+          @Override
+          public StatusCode.Code getCode() {
+            return StatusCode.Code.FAILED_PRECONDITION;
+          }
+
+          @Override
+          public @Nullable Object getTransportCode() {
+            return null;
+          }
+        },
+        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 @@
         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();