feat(gax): add RewindableStreamBuffer for single-chunk rewinds and seeks

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

@@ -0,0 +1,176 @@
+/*
+ * 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.resumable;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import com.google.common.io.ByteStreams;
+import com.google.errorprone.annotations.concurrent.GuardedBy;
+import com.google.protobuf.ByteString;
+import java.io.IOException;
+import java.io.InputStream;
+import org.jspecify.annotations.NullMarked;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * A package-private stream buffer that supports single-chunk rewind and seeking over an {@link
+ * InputStream} for resumable uploads.
+ */
+@NullMarked
+class RewindableStreamBuffer implements AutoCloseable {
+
+  private final Object lock = new Object();
+  private final InputStream source;
+
+  @GuardedBy("lock")
+  private @Nullable ByteString currentChunk;
+
+  @GuardedBy("lock")
+  private long currentChunkStartOffset = 0L;
+
+  @GuardedBy("lock")
+  private long streamPosition = 0L;
+
+  @GuardedBy("lock")
+  private boolean endOfStream = false;
+
+  RewindableStreamBuffer(InputStream source) {
+    this.source = checkNotNull(source);
+  }
+
+  /**
+   * Reads up to {@code maxBytes} from the underlying stream starting at {@code targetOffset}. If
+   * bytes for this offset are already in the buffer (e.g. during a retry attempt), returns the
+   * uncommitted slice without reading anew from the source stream.
+   *
+   * @param targetOffset expected start offset of the chunk
+   * @param maxBytes maximum bytes to return in this chunk
+   * @return byte data for the chunk, or {@link ByteString#EMPTY} if EOF is reached
+   * @throws IOException on I/O error reading from stream
+   */
+  ByteString readChunk(long targetOffset, int maxBytes) throws IOException {
+    checkArgument(maxBytes > 0, "maxBytes must be > 0");
+
+    synchronized (lock) {
+      if (currentChunk != null
+          && targetOffset >= currentChunkStartOffset
+          && targetOffset < currentChunkStartOffset + currentChunk.size()) {
+        int offsetInChunk = (int) (targetOffset - currentChunkStartOffset);
+        int len = Math.min(maxBytes, currentChunk.size() - offsetInChunk);
+        return currentChunk.substring(offsetInChunk, offsetInChunk + len);
+      }
+
+      if (endOfStream) {
+        return ByteString.EMPTY;
+      }
+
+      checkArgument(
+          targetOffset == streamPosition,
+          "targetOffset (%s) does not match streamPosition (%s)",
+          targetOffset,
+          streamPosition);
+
+      byte[] buffer = new byte[maxBytes];
+      int totalRead = 0;
+      while (totalRead < maxBytes) {
+        int read = source.read(buffer, totalRead, maxBytes - totalRead);
+        if (read == -1) {
+          endOfStream = true;
+          break;
+        }
+        totalRead += read;
+      }
+
+      if (totalRead == 0) {
+        return ByteString.EMPTY;
+      }
+
+      streamPosition += totalRead;
+      currentChunk = ByteString.copyFrom(buffer, 0, totalRead);
+      currentChunkStartOffset = targetOffset;
+      return currentChunk;
+    }
+  }
+
+  /** Returns true if the end of the underlying stream has been reached. */
+  boolean isEndOfStream() {
+    synchronized (lock) {
+      return endOfStream;
+    }
+  }
+
+  /**
+   * Seeks the buffer position to {@code committedOffset} as reported by {@code queryStatus}.
+   *
+   * @param committedOffset server's committed byte count
+   * @throws IOException on I/O error skipping bytes in the source stream
+   */
+  void seek(long committedOffset) throws IOException {
+    synchronized (lock) {
+      if (committedOffset < currentChunkStartOffset) {
+        throw new IllegalArgumentException(
+            "Cannot seek backwards before current chunk start offset " + currentChunkStartOffset);
+      }
+      if (currentChunk != null && committedOffset < currentChunkStartOffset + currentChunk.size()) {
+        return;
+      }
+      currentChunk = null;
+      long bytesToSkip = committedOffset - streamPosition;
+      if (bytesToSkip > 0) {
+        ByteStreams.skipFully(source, bytesToSkip);
+        streamPosition = committedOffset;
+      }
+      currentChunkStartOffset = committedOffset;
+    }
+  }
+
+  /**
+   * Commits and discards buffered data up to {@code committedOffset}.
+   *
+   * @param committedOffset newly confirmed committed offset
+   */
+  void commit(long committedOffset) {
+    synchronized (lock) {
+      if (currentChunk != null
+          && committedOffset >= currentChunkStartOffset + currentChunk.size()) {
+        currentChunk = null;
+        currentChunkStartOffset = committedOffset;
+      }
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    synchronized (lock) {
+      source.close();
+    }
+  }
+}

sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/resumable/RewindableStreamBufferTest.java

@@ -0,0 +1,187 @@
+/*
+ * 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.resumable;
+
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import org.junit.jupiter.api.Test;
+
+class RewindableStreamBufferTest {
+
+  @Test
+  void testReadSequentialChunks() throws IOException {
+    byte[] data = "HelloWorld123456".getBytes(StandardCharsets.UTF_8); // 16 bytes
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    // Read first chunk (8 bytes: "HelloWor")
+    ByteString chunk1 = buffer.readChunk(0, 8);
+    assertThat(chunk1.toStringUtf8()).isEqualTo("HelloWor");
+    assertThat(buffer.isEndOfStream()).isFalse();
+
+    // Commit first chunk
+    buffer.commit(8);
+
+    // Read second chunk (8 bytes: "ld123456")
+    ByteString chunk2 = buffer.readChunk(8, 8);
+    assertThat(chunk2.toStringUtf8()).isEqualTo("ld123456");
+
+    // Read at EOF
+    buffer.commit(16);
+    ByteString chunk3 = buffer.readChunk(16, 8);
+    assertThat(chunk3.isEmpty()).isTrue();
+    assertThat(buffer.isEndOfStream()).isTrue();
+
+    buffer.close();
+  }
+
+  @Test
+  void testSeekAndRewindWithinChunk() throws IOException {
+    byte[] data = "0123456789ABCDEF".getBytes(StandardCharsets.UTF_8); // 16 bytes
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    // Read first chunk of 8 bytes ("01234567")
+    ByteString chunk1 = buffer.readChunk(0, 8);
+    assertThat(chunk1.toStringUtf8()).isEqualTo("01234567");
+
+    // Simulate failure where server only committed 3 bytes ("012")
+    buffer.seek(3);
+
+    // Retry reading from offset 3 (should return slice "34567")
+    ByteString retrySlice = buffer.readChunk(3, 8);
+    assertThat(retrySlice.toStringUtf8()).isEqualTo("34567");
+
+    // Once server commits full 8 bytes
+    buffer.commit(8);
+
+    // Read next chunk ("89ABCDEF")
+    ByteString chunk2 = buffer.readChunk(8, 8);
+    assertThat(chunk2.toStringUtf8()).isEqualTo("89ABCDEF");
+
+    buffer.close();
+  }
+
+  @Test
+  void testPartialReadSmallerThanChunkSize() throws IOException {
+    byte[] data = "Small".getBytes(StandardCharsets.UTF_8); // 5 bytes
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    ByteString chunk = buffer.readChunk(0, 10);
+    assertThat(chunk.toStringUtf8()).isEqualTo("Small");
+    assertThat(buffer.isEndOfStream()).isTrue(); // Encountered EOF while reading
+
+    ByteString nextChunk = buffer.readChunk(5, 10);
+    assertThat(nextChunk.isEmpty()).isTrue();
+    assertThat(buffer.isEndOfStream()).isTrue();
+
+    buffer.close();
+  }
+
+  @Test
+  void testEmptyStream() throws IOException {
+    ByteArrayInputStream stream = new ByteArrayInputStream(new byte[0]);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    ByteString chunk = buffer.readChunk(0, 8);
+    assertThat(chunk.isEmpty()).isTrue();
+    assertThat(buffer.isEndOfStream()).isTrue();
+
+    buffer.close();
+  }
+
+  @Test
+  void testSeekFromBeginning() throws IOException {
+    byte[] data = "0123456789ABCDEF".getBytes(StandardCharsets.UTF_8); // 16 bytes
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    // Seek directly to offset 8 before any reads (e.g. resuming session)
+    buffer.seek(8);
+
+    ByteString chunk = buffer.readChunk(8, 8);
+    assertThat(chunk.toStringUtf8()).isEqualTo("89ABCDEF");
+    assertThat(buffer.isEndOfStream()).isFalse();
+
+    buffer.close();
+  }
+
+  @Test
+  void testSeekForwardAcrossChunks() throws IOException {
+    byte[] data = "0123456789ABCDEF".getBytes(StandardCharsets.UTF_8); // 16 bytes
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    ByteString chunk1 = buffer.readChunk(0, 4);
+    assertThat(chunk1.toStringUtf8()).isEqualTo("0123");
+
+    // Seek past chunk 1 to offset 10
+    buffer.seek(10);
+
+    ByteString chunk2 = buffer.readChunk(10, 4);
+    assertThat(chunk2.toStringUtf8()).isEqualTo("ABCD");
+
+    buffer.close();
+  }
+
+  @Test
+  void testInvalidMaxBytes() {
+    ByteArrayInputStream stream = new ByteArrayInputStream(new byte[0]);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+    assertThrows(IllegalArgumentException.class, () -> buffer.readChunk(0, 0));
+    assertThrows(IllegalArgumentException.class, () -> buffer.readChunk(0, -1));
+  }
+
+  @Test
+  void testReadChunkMismatchedOffsetThrows() {
+    byte[] data = "0123456789ABCDEF".getBytes(StandardCharsets.UTF_8);
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    // streamPosition is 0, but passing targetOffset 5 without seeking
+    assertThrows(IllegalArgumentException.class, () -> buffer.readChunk(5, 4));
+  }
+
+  @Test
+  void testCannotSeekBackwardsBeforeChunkStart() throws IOException {
+    byte[] data = "0123456789ABCDEF".getBytes(StandardCharsets.UTF_8);
+    ByteArrayInputStream stream = new ByteArrayInputStream(data);
+    RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream);
+
+    buffer.seek(8);
+    assertThrows(IllegalArgumentException.class, () -> buffer.seek(4));
+  }
+}