From 6e8ee1a2a62cdbae1699a2a8db03c9b57b372a34 Mon Sep 17 00:00:00 2001 From: whowes Date: Thu, 20 Aug 2026 22:54:39 +0000 Subject: [PATCH] feat(gax): add RewindableStreamBuffer for single-chunk rewinds and seeks --- .../gax/resumable/RewindableStreamBuffer.java | 176 +++++++++++++++++ .../resumable/RewindableStreamBufferTest.java | 186 ++++++++++++++++++ 2 files changed, 362 insertions(+) create mode 100644 sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/resumable/RewindableStreamBuffer.java create mode 100644 sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/resumable/RewindableStreamBufferTest.java diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/resumable/RewindableStreamBuffer.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/resumable/RewindableStreamBuffer.java new file mode 100644 index 000000000000..f3f97e2e0005 --- /dev/null +++ b/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 java.io.IOException; +import java.io.InputStream; +import java.util.Arrays; +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 byte @Nullable [] 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 an empty array if EOF is reached + * @throws IOException on I/O error reading from stream + */ + byte[] readChunk(long targetOffset, int maxBytes) throws IOException { + checkArgument(maxBytes > 0, "maxBytes must be > 0"); + + synchronized (lock) { + if (currentChunk != null + && targetOffset >= currentChunkStartOffset + && targetOffset < currentChunkStartOffset + currentChunk.length) { + int offsetInChunk = (int) (targetOffset - currentChunkStartOffset); + int len = Math.min(maxBytes, currentChunk.length - offsetInChunk); + return Arrays.copyOfRange(currentChunk, offsetInChunk, offsetInChunk + len); + } + + if (endOfStream) { + return new byte[0]; + } + + 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 new byte[0]; + } + + streamPosition += totalRead; + currentChunk = totalRead == maxBytes ? buffer : Arrays.copyOf(buffer, 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.length) { + 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.length) { + currentChunk = null; + currentChunkStartOffset = committedOffset; + } + } + } + + @Override + public void close() throws IOException { + synchronized (lock) { + source.close(); + } + } +} diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/resumable/RewindableStreamBufferTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/resumable/RewindableStreamBufferTest.java new file mode 100644 index 000000000000..17751b827729 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/resumable/RewindableStreamBufferTest.java @@ -0,0 +1,186 @@ +/* + * 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 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") + byte[] chunk1 = buffer.readChunk(0, 8); + assertThat(new String(chunk1, StandardCharsets.UTF_8)).isEqualTo("HelloWor"); + assertThat(buffer.isEndOfStream()).isFalse(); + + // Commit first chunk + buffer.commit(8); + + // Read second chunk (8 bytes: "ld123456") + byte[] chunk2 = buffer.readChunk(8, 8); + assertThat(new String(chunk2, StandardCharsets.UTF_8)).isEqualTo("ld123456"); + + // Read at EOF + buffer.commit(16); + byte[] chunk3 = buffer.readChunk(16, 8); + assertThat(chunk3).isEmpty(); + 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") + byte[] chunk1 = buffer.readChunk(0, 8); + assertThat(new String(chunk1, StandardCharsets.UTF_8)).isEqualTo("01234567"); + + // Simulate failure where server only committed 3 bytes ("012") + buffer.seek(3); + + // Retry reading from offset 3 (should return slice "34567") + byte[] retrySlice = buffer.readChunk(3, 8); + assertThat(new String(retrySlice, StandardCharsets.UTF_8)).isEqualTo("34567"); + + // Once server commits full 8 bytes + buffer.commit(8); + + // Read next chunk ("89ABCDEF") + byte[] chunk2 = buffer.readChunk(8, 8); + assertThat(new String(chunk2, StandardCharsets.UTF_8)).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); + + byte[] chunk = buffer.readChunk(0, 10); + assertThat(new String(chunk, StandardCharsets.UTF_8)).isEqualTo("Small"); + assertThat(buffer.isEndOfStream()).isTrue(); // Encountered EOF while reading + + byte[] nextChunk = buffer.readChunk(5, 10); + assertThat(nextChunk).isEmpty(); + assertThat(buffer.isEndOfStream()).isTrue(); + + buffer.close(); + } + + @Test + void testEmptyStream() throws IOException { + ByteArrayInputStream stream = new ByteArrayInputStream(new byte[0]); + RewindableStreamBuffer buffer = new RewindableStreamBuffer(stream); + + byte[] chunk = buffer.readChunk(0, 8); + assertThat(chunk).isEmpty(); + 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); + + byte[] chunk = buffer.readChunk(8, 8); + assertThat(new String(chunk, StandardCharsets.UTF_8)).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); + + byte[] chunk1 = buffer.readChunk(0, 4); + assertThat(new String(chunk1, StandardCharsets.UTF_8)).isEqualTo("0123"); + + // Seek past chunk 1 to offset 10 + buffer.seek(10); + + byte[] chunk2 = buffer.readChunk(10, 4); + assertThat(new String(chunk2, StandardCharsets.UTF_8)).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)); + } +}