feat(gax): add chunk upload recovery loop for resumable uploads

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;
@@ -78,6 +80,9 @@ public class ResumableUploadCallableImpl<RequestT, ResponseT>
   private final ClientContext clientContext;
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       retryingUploadChunkCallable;
+  private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
+      retryingQueryCallable;
+  private final ExponentialRetryAlgorithm recoveryAlgorithm;
 
   public ResumableUploadCallableImpl(
       ResumableUploadClient<RequestT, ResponseT> client,
@@ -90,6 +95,11 @@ public ResumableUploadCallableImpl(
     this.retryingUploadChunkCallable =
         createRetryingCallable(
             client.uploadChunkCallable(), ResumableUploadCommand.UPLOAD, clientContext);
+    this.retryingQueryCallable =
+        createRetryingCallable(
+            client.queryStatusCallable(), ResumableUploadCommand.QUERY, clientContext);
+    this.recoveryAlgorithm =
+        new ExponentialRetryAlgorithm(RETRY_SETTINGS, clientContext.getClock());
   }
 
   @Override
@@ -113,9 +123,11 @@ public ResumableUploadFuture<ResponseT> futureCall(
     return ResumableUploadFutureImpl.create(
         startFuture,
         retryingUploadChunkCallable,
+        retryingQueryCallable,
         payload,
         effectiveSettings,
-        clientContext.getDefaultCallContext());
+        clientContext,
+        recoveryAlgorithm);
   }
 
   @Override

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

@@ -38,10 +38,19 @@
 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.ExponentialRetryAlgorithm;
+import com.google.api.gax.retrying.TimedAttemptSettings;
+import com.google.api.gax.rpc.ResumableUploadErrorClassifier.Category;
 import com.google.common.util.concurrent.MoreExecutors;
+import java.io.IOException;
 import java.io.InputStream;
 import java.util.concurrent.CancellationException;
+import java.util.concurrent.Future;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -56,23 +65,37 @@
 
   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 ExponentialRetryAlgorithm recoveryAlgorithm;
+  private final ScheduledExecutorService executor;
   private final SettableApiFuture<ResponseT> uploadResultFuture = SettableApiFuture.create();
-  private volatile @Nullable ApiFuture<?> inFlightFuture;
+  private volatile @Nullable Future<?> inFlightFuture;
+  private TimedAttemptSettings recoverySettings;
+  private boolean madeProgressSinceRecovery = true;
 
   ResumableUploadChunkCoordinator(
       UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>> uploadChunkCallable,
+      UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>> queryStatusCallable,
       String uploadUrl,
       InputStream payload,
       int chunkSize,
-      ApiCallContext callContext) {
+      ApiCallContext callContext,
+      ExponentialRetryAlgorithm recoveryAlgorithm,
+      ScheduledExecutorService executor) {
     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");
+    this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null");
+    this.executor = checkNotNull(executor, "executor must not be null");
+    this.recoverySettings = recoveryAlgorithm.createFirstAttempt();
     this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl);
   }
 
@@ -83,42 +106,35 @@
   void start() {
     uploadResultFuture.addListener(
         () -> {
-          ApiFuture<?> inFlight = inFlightFuture;
+          Future<?> inFlight = inFlightFuture;
           if (uploadResultFuture.isCancelled() && inFlight != null) {
             inFlight.cancel(true);
           }
         },
         MoreExecutors.directExecutor());
-    transmitChunk();
+    try {
+      buffer.fill();
+      transmitChunk();
+    } catch (Throwable t) {
+      uploadResultFuture.setException(t);
+    }
   }
 
   private void transmitChunk() {
     try {
       // Abort if the session was already completed or canceled.
       if (uploadResultFuture.isDone()) {
         return;
       }
 
-      // Read the next chunk slice from the payload stream.
-      buffer.fill();
-
-      // Determine if this is the final chunk and build the chunk request.
-      ChunkUploadRequest chunkRequest =
-          ChunkUploadRequest.newBuilder()
-              .setUploadUrl(uploadUrl)
-              .setPayload(buffer.getPayload())
-              .setOffset(buffer.getBufferBaseOffset())
-              .setFinal(buffer.isFinal())
-              .build();
+      ChunkUploadRequest chunkRequest = buildCurrentChunkRequest();
 
       // Dispatch the chunk upload call and register the in-flight future for cancellation.
-      boolean isFinal = chunkRequest.isFinal();
       ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
           uploadChunkCallable.futureCall(chunkRequest, callContext);
       if (!tryRegisterInFlightFuture(chunkFuture)) {
         return;
       }
-
       ApiFutures.addCallback(
           chunkFuture,
           new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
@@ -127,15 +143,88 @@
               if (uploadResultFuture.isDone()) {
                 return;
               }
-              if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
-                uploadResultFuture.set(response.getResponse());
-              } else if (isFinal) {
-                uploadResultFuture.setException(
+              if (response.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+                recover(
                     new IllegalStateException(
-                        "Upload stream ended and final chunk was transmitted, but server returned"
-                            + " incomplete status"));
+                        "Chunk upload response missing X-Goog-Upload-Status header for upload URL: "
+                            + uploadUrl));
               } else {
-                transmitChunk();
+                try {
+                  handleChunkResponse(response);
+                } catch (Throwable t) {
+                  uploadResultFuture.setException(t);
+                }
+              }
+            }
+
+            @Override
+            public void onFailure(Throwable t) {
+              if (t instanceof CancellationException || uploadResultFuture.isDone()) {
+                return;
+              }
+              Category category =
+                  ResumableUploadErrorClassifier.classify(t, ResumableUploadCommand.UPLOAD);
+              if (category == Category.RECOVERABLE) {
+                recover(t);
+              } else {
+                // Category.TRANSIENT errors reaching here have already exhausted their retry budget
+                // in the underlying RetryingCallable and become fatal per protocol specification.
+                uploadResultFuture.setException(t);
+              }
+            }
+          },
+          MoreExecutors.directExecutor());
+    } catch (Throwable t) {
+      uploadResultFuture.setException(t);
+    }
+  }
+
+  private void recover(Throwable cause) {
+    try {
+      if (madeProgressSinceRecovery) {
+        recoverySettings =
+            recoveryAlgorithm.createNextAttempt(recoveryAlgorithm.createFirstAttempt());
+      } else {
+        recoverySettings = recoveryAlgorithm.createNextAttempt(recoverySettings);
+      }
+      madeProgressSinceRecovery = false;
+      if (!recoveryAlgorithm.shouldRetry(recoverySettings)) {
+        uploadResultFuture.setException(cause);
+        return;
+      }
+      tryRegisterInFlightFuture(
+          executor.schedule(
+              this::queryStatus,
+              recoverySettings.getRandomizedRetryDelayDuration().toNanos(),
+              TimeUnit.NANOSECONDS));
+    } catch (Throwable t) {
+      uploadResultFuture.setException(t);
+    }
+  }
+
+  private void queryStatus() {
+    try {
+      if (uploadResultFuture.isDone()) {
+        return;
+      }
+      // Dispatch the query status call and register the in-flight future for cancellation.
+      ApiFuture<QueryStatusResponse<ResponseT>> queryFuture =
+          queryStatusCallable.futureCall(QueryStatusRequest.create(uploadUrl), callContext);
+      if (!tryRegisterInFlightFuture(queryFuture)) {
+        return;
+      }
+      ApiFutures.addCallback(
+          queryFuture,
+          new ApiFutureCallback<QueryStatusResponse<ResponseT>>() {
+            @Override
+            public void onSuccess(QueryStatusResponse<ResponseT> queryResponse) {
+              if (uploadResultFuture.isDone()) {
+                return;
+              }
+              try {
+                handleQueryResponse(queryResponse);
+              } catch (Throwable t) {
+                uploadResultFuture.setException(t);
               }
             }
 
@@ -157,12 +246,58 @@
    * Registers the in-flight future for possible cancellation, returning false if the upload was
    * already cancelled.
    */
-  private boolean tryRegisterInFlightFuture(ApiFuture<?> future) {
+  private boolean tryRegisterInFlightFuture(Future<?> future) {
     this.inFlightFuture = future;
     if (uploadResultFuture.isCancelled()) {
       future.cancel(true);
       return false;
     }
     return true;
   }
+
+  private void handleQueryResponse(QueryStatusResponse<ResponseT> queryResponse)
+      throws IOException {
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.UNKNOWN) {
+      throw new IllegalStateException(
+          "Query status response missing X-Goog-Upload-Status header for upload URL: " + uploadUrl);
+    }
+    if (queryResponse.getUploadStatus() == ResumableUploadStatus.FINAL) {
+      uploadResultFuture.set(queryResponse.getResponse());
+      return;
+    }
+    Long committedOffset = queryResponse.getCommittedOffset();
+    if (committedOffset == null) {
+      throw new IllegalStateException(
+          "Incomplete query status response did not include a committed offset for upload URL: "
+              + uploadUrl);
+    }
+    buffer.realignTo(committedOffset);
+    transmitChunk();
+  }
+
+  private void handleChunkResponse(ChunkUploadResponse<ResponseT> response) throws IOException {
+    if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
+      uploadResultFuture.set(response.getResponse());
+    } else if (buffer.isFinal()) {
+      uploadResultFuture.setException(
+          new IllegalStateException(
+              "Upload stream ended and final chunk was transmitted, but server returned"
+                  + " incomplete status for upload URL: "
+                  + uploadUrl));
+    } else {
+      madeProgressSinceRecovery = true;
+      buffer.fill();
+      transmitChunk();
+    }
+  }
+
+  private ChunkUploadRequest buildCurrentChunkRequest() {
+    // Determine if this is the final chunk and build the chunk request.
+    return ChunkUploadRequest.newBuilder()
+        .setUploadUrl(uploadUrl)
+        .setPayload(buffer.getPayload())
+        .setOffset(buffer.getBufferBaseOffset())
+        .setFinal(buffer.isFinal())
+        .build();
+  }
 }