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

@@ -209,24 +209,22 @@ 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);
+    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())
@@ -236,21 +234,20 @@ private ChunkUploadRequest buildCurrentChunkRequest() {
         .build();
   }
 
-  private static FailedPreconditionException protocolViolation(String message) {
-    return new FailedPreconditionException(
-        message,
-        null,
-        new StatusCode() {
-          @Override
-          public StatusCode.Code getCode() {
-            return StatusCode.Code.FAILED_PRECONDITION;
-          }
+  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;
-          }
-        },
-        false);
+        @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();