feat(gax): add rewindable stream buffer for chunk recovery

Introduce RewindableStreamBuffer managing a single-chunk buffer over an
InputStream, supporting forward compaction and topping up upon recovery
realignment without mark()/reset(). Enforces boundaries by throwing
FailedPreconditionException when a server offset is below the base
offset or beyond the current buffer window.

Use payloadLength in ChunkUploadRequest to avoid allocating temporary
byte arrays for full-sized chunks while reusing a single backing array.

Author: whowes

sdk-platform-java/gax-java/gax-httpjson/src/main/java/com/google/api/gax/httpjson/ResumableUploadChunkCallable.java

@@ -46,6 +46,7 @@
 import java.io.ByteArrayInputStream;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
@@ -81,7 +82,12 @@ public Map<String, List<String>> getQueryParamNames(ChunkUploadRequest request)
 
                 @Override
                 public byte[] getBinaryRequestBody(ChunkUploadRequest request) {
-                  return request.getPayload();
+                  int length = request.getPayloadLength();
+                  byte[] payload = request.getPayload();
+                  if (length == payload.length) {
+                    return payload;
+                  }
+                  return Arrays.copyOf(payload, length);
                 }
 
                 @Override
@@ -110,7 +116,7 @@ private ResumableUploadChunkCallable(
   public ApiFuture<ChunkUploadResponse<ResponseT>> futureCall(
       ChunkUploadRequest request, @Nullable ApiCallContext inputContext) {
     Preconditions.checkNotNull(request);
-    boolean isPayloadEmpty = request.getPayload().length == 0;
+    boolean isPayloadEmpty = request.getPayloadLength() == 0;
     String command;
     if (request.isFinal()) {
       command = !isPayloadEmpty ? COMMAND_UPLOAD_FINALIZE : COMMAND_FINALIZE;

sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/resumable/ChunkUploadRequest.java

@@ -32,6 +32,7 @@
 import com.google.api.core.BetaApi;
 import com.google.api.core.InternalApi;
 import com.google.auto.value.AutoValue;
+import com.google.common.base.Preconditions;
 import org.jspecify.annotations.NullMarked;
 
 /** Request value object for uploading a chunk to an active resumable upload session. */
@@ -48,6 +49,9 @@ public abstract class ChunkUploadRequest {
   @SuppressWarnings("mutable")
   public abstract byte[] getPayload();
 
+  /** The number of bytes within {@link #getPayload()} to upload. */
+  public abstract int getPayloadLength();
+
   /** The byte offset of this chunk in the overall stream. */
   public abstract long getOffset();
 
@@ -56,8 +60,12 @@ public abstract class ChunkUploadRequest {
 
   public abstract Builder toBuilder();
 
+  private static final int UNSET_PAYLOAD_LENGTH = Integer.MIN_VALUE;
+
   public static Builder newBuilder() {
-    return new AutoValue_ChunkUploadRequest.Builder().setFinal(false);
+    return new AutoValue_ChunkUploadRequest.Builder()
+        .setFinal(false)
+        .setPayloadLength(UNSET_PAYLOAD_LENGTH);
   }
 
   @AutoValue.Builder
@@ -66,10 +74,29 @@ public abstract static class Builder {
 
     public abstract Builder setPayload(byte[] payload);
 
+    public abstract Builder setPayloadLength(int payloadLength);
+
     public abstract Builder setOffset(long offset);
 
     public abstract Builder setFinal(boolean isFinal);
 
-    public abstract ChunkUploadRequest build();
+    abstract byte[] getPayload();
+
+    abstract int getPayloadLength();
+
+    abstract ChunkUploadRequest autoBuild();
+
+    public ChunkUploadRequest build() {
+      if (getPayloadLength() == UNSET_PAYLOAD_LENGTH) {
+        setPayloadLength(getPayload().length);
+      }
+      ChunkUploadRequest request = autoBuild();
+      Preconditions.checkArgument(
+          request.getPayloadLength() >= 0, "payloadLength must be non-negative");
+      Preconditions.checkArgument(
+          request.getPayloadLength() <= request.getPayload().length,
+          "payloadLength exceeds payload array length");
+      return request;
+    }
   }
 }

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

@@ -38,12 +38,13 @@
 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.common.io.ByteStreams;
+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.Arrays;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.Executor;
 import org.jspecify.annotations.NullMarked;
@@ -52,42 +53,45 @@
 /**
  * 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> {
 
-  private static final byte[] EMPTY_PAYLOAD = new byte[0];
-
   // Serializes buffer mutations across multiple threads (i.e. from retry/recovery)
   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 InputStream payload;
-  private final byte[] buffer;
-  private final int chunkSize;
+  private final RewindableStreamBuffer buffer;
   private final ApiCallContext callContext;
   private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
   private volatile @Nullable ApiFuture<?> currentChunkFuture;
 
   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");
-    this.payload = checkNotNull(payload, "payload must not be null");
-    this.chunkSize = chunkSize;
+    checkNotNull(payload, "payload must not be null");
     checkNotNull(clientContext, "clientContext must not be null");
     this.callContext = clientContext.getDefaultCallContext();
-    this.buffer = new byte[chunkSize];
+    this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl);
   }
 
   ApiFuture<ResponseT> start() {
@@ -104,75 +108,39 @@ 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.
-    int bytesRead;
     try {
-      bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize);
-    } catch (IOException e) {
-      result.setException(e);
-      return;
+      buffer.fill(currentOffset);
+      dispatchCurrentChunk();
+    } catch (Throwable t) {
+      result.setException(t);
     }
+  }
 
-    // Determine if this is the final chunk and build the chunk request.
-    boolean isFinal = bytesRead < chunkSize;
-    byte[] chunkPayload;
-    if (bytesRead == chunkSize) {
-      chunkPayload = buffer;
-    } else if (bytesRead == 0) {
-      chunkPayload = EMPTY_PAYLOAD;
-    } else {
-      chunkPayload = Arrays.copyOf(buffer, bytesRead);
+  private void dispatchCurrentChunk() {
+    if (result.isDone()) {
+      return;
     }
-
-    ChunkUploadRequest chunkRequest =
-        ChunkUploadRequest.newBuilder()
-            .setUploadUrl(uploadUrl)
-            .setPayload(chunkPayload)
-            .setOffset(currentOffset)
-            .setFinal(isFinal)
-            .build();
-
-    // Dispatch the chunk upload call and register the in-flight future for cancellation.
-    long chunkLength = chunkPayload.length;
     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);
@@ -183,4 +151,107 @@ public void onFailure(Throwable t) {
       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);
+  }
 }