feat(gax): add RewindableStreamBuffer in prep for chunk upload recovery

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

@@ -39,10 +39,8 @@
 import com.google.api.gax.resumable.ChunkUploadRequest;
 import com.google.api.gax.resumable.ChunkUploadResponse;
 import com.google.api.gax.resumable.ResumableUploadStatus;
-import com.google.common.io.ByteStreams;
 import com.google.common.util.concurrent.MoreExecutors;
 import java.io.InputStream;
-import java.util.Arrays;
 import java.util.concurrent.CancellationException;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
@@ -56,14 +54,10 @@
 @NullMarked
 final class ResumableUploadChunkCoordinator<ResponseT> {
 
-  private static final byte[] EMPTY_PAYLOAD = new byte[0];
-
   private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
       uploadChunkCallable;
   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> uploadResultFuture = SettableApiFuture.create();
   private volatile @Nullable ApiFuture<?> inFlightFuture;
@@ -77,10 +71,9 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
     this.uploadChunkCallable =
         checkNotNull(uploadChunkCallable, "uploadChunkCallable 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");
     this.callContext = checkNotNull(callContext, "callContext must not be null");
-    this.buffer = new byte[chunkSize];
+    this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl);
   }
 
   ApiFuture<ResponseT> getFuture() {
@@ -96,40 +89,30 @@ void start() {
           }
         },
         MoreExecutors.directExecutor());
-    transmitChunk(0L);
+    transmitChunk();
   }
 
-  private void transmitChunk(long currentOffset) {
+  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.
-      int bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize);
+      buffer.fill();
 
       // 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);
-      }
-
       ChunkUploadRequest chunkRequest =
           ChunkUploadRequest.newBuilder()
               .setUploadUrl(uploadUrl)
-              .setPayload(chunkPayload)
-              .setOffset(currentOffset)
-              .setFinal(isFinal)
+              .setPayload(buffer.getPayload())
+              .setOffset(buffer.getBufferBaseOffset())
+              .setFinal(buffer.isFinal())
               .build();
 
       // Dispatch the chunk upload call and register the in-flight future for cancellation.
-      long chunkLength = chunkPayload.length;
+      boolean isFinal = chunkRequest.isFinal();
       ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
           uploadChunkCallable.futureCall(chunkRequest, callContext);
       if (!tryRegisterInFlightFuture(chunkFuture)) {
@@ -144,7 +127,6 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
               if (uploadResultFuture.isDone()) {
                 return;
               }
-              long nextOffset = currentOffset + chunkLength;
               if (response.getUploadStatus() == ResumableUploadStatus.FINAL) {
                 uploadResultFuture.set(response.getResponse());
               } else if (isFinal) {
@@ -153,7 +135,7 @@ public void onSuccess(ChunkUploadResponse<ResponseT> response) {
                         "Upload stream ended and final chunk was transmitted, but server returned"
                             + " incomplete status"));
               } else {
-                transmitChunk(nextOffset);
+                transmitChunk();
               }
             }
 

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

@@ -0,0 +1,150 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are
+ * met:
+ *
+ *     * Redistributions of source code must retain the above copyright
+ * notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above
+ * copyright notice, this list of conditions and the following disclaimer
+ * in the documentation and/or other materials provided with the
+ * distribution.
+ *     * Neither the name of Google LLC nor the names of its
+ * contributors may be used to endorse or promote products derived from
+ * this software without specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+ * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+ * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+ * A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+ * OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+ * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
+ * LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
+ * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
+ * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+ * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ */
+package com.google.api.gax.rpc;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.common.io.ByteStreams;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Arrays;
+import org.jspecify.annotations.NullMarked;
+
+/** Manages a single-chunk rewindable buffer over an {@link InputStream} for resumable uploads. */
+@NullMarked
+final class RewindableStreamBuffer {
+
+  private static final byte[] EMPTY_PAYLOAD = new byte[0];
+
+  private final InputStream inputStream;
+  private final int chunkSize;
+  private final String uploadUrl;
+  private final byte[] buffer;
+
+  private long bufferBaseOffset;
+  private int payloadLength;
+  private boolean isFinal;
+  private boolean streamExhausted;
+
+  RewindableStreamBuffer(InputStream inputStream, int chunkSize, String uploadUrl) {
+    this.inputStream = checkNotNull(inputStream, "inputStream must not be null");
+    checkArgument(chunkSize > 0, "chunkSize must be > 0");
+    this.chunkSize = chunkSize;
+    this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
+    this.buffer = new byte[chunkSize];
+    this.bufferBaseOffset = 0L;
+    this.payloadLength = 0;
+    this.isFinal = false;
+    this.streamExhausted = false;
+  }
+
+  /**
+   * Advances the buffer past the current payload, reading up to chunk size from the stream.
+   *
+   * @throws IOException if reading from the stream fails
+   */
+  void fill() throws IOException {
+    this.bufferBaseOffset += payloadLength;
+    this.payloadLength = ByteStreams.read(inputStream, buffer, 0, chunkSize);
+    this.isFinal = (payloadLength < chunkSize);
+    if (this.isFinal) {
+      this.streamExhausted = true;
+    }
+  }
+
+  /**
+   * Realigns the buffer window to {@code committedOffset}.
+   *
+   * <p>Compacts forward within the existing buffer to discard already-committed bytes, and then
+   * tops up the buffer to capacity from the underlying stream.
+   *
+   * @param committedOffset the server's committed byte offset
+   * @throws IllegalStateException if {@code committedOffset} is below the buffer's base offset or
+   *     beyond the current buffer window
+   * @throws IOException if reading from the stream fails
+   */
+  void realignTo(long committedOffset) throws IOException {
+    if (committedOffset < bufferBaseOffset) {
+      throw new IllegalStateException(
+          String.format(
+              "Server committed offset %d is below buffer base offset %d for upload URL %s; cannot"
+                  + " rewind stream before buffer base",
+              committedOffset, bufferBaseOffset, uploadUrl));
+    }
+
+    if (committedOffset > bufferBaseOffset + payloadLength) {
+      throw new IllegalStateException(
+          String.format(
+              "Server committed offset %d is beyond current buffer window [%d, %d] for upload URL"
+                  + " %s",
+              committedOffset, bufferBaseOffset, bufferBaseOffset + payloadLength, uploadUrl));
+    }
+
+    int committedWithinBuffer = (int) (committedOffset - bufferBaseOffset);
+    int remainingBytes = payloadLength - committedWithinBuffer;
+
+    if (remainingBytes > 0 && committedWithinBuffer > 0) {
+      System.arraycopy(buffer, committedWithinBuffer, buffer, 0, remainingBytes);
+    }
+
+    this.bufferBaseOffset = committedOffset;
+    this.payloadLength = remainingBytes;
+
+    if (!streamExhausted && payloadLength < chunkSize) {
+      int space = chunkSize - payloadLength;
+      int additionalRead = ByteStreams.read(inputStream, buffer, payloadLength, space);
+      payloadLength += additionalRead;
+      if (additionalRead < space) {
+        streamExhausted = true;
+      }
+    }
+
+    this.isFinal = streamExhausted;
+  }
+
+  byte[] getPayload() {
+    if (payloadLength == buffer.length) {
+      return buffer;
+    }
+    if (payloadLength == 0) {
+      return EMPTY_PAYLOAD;
+    }
+    return Arrays.copyOf(buffer, payloadLength);
+  }
+
+  long getBufferBaseOffset() {
+    return bufferBaseOffset;
+  }
+
+  boolean isFinal() {
+    return isFinal;
+  }
+}