From 7431c49af5d8eb120bef637950d3f4775cdc6b90 Mon Sep 17 00:00:00 2001 From: whowes Date: Sat, 12 Sep 2026 15:44:12 +0000 Subject: [PATCH] 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 IllegalStateException when a server offset is below the base offset or beyond the current buffer window. --- .../rpc/ResumableUploadChunkCoordinator.java | 36 +-- .../api/gax/rpc/RewindableStreamBuffer.java | 155 +++++++++++++ .../gax/rpc/RewindableStreamBufferTest.java | 217 ++++++++++++++++++ 3 files changed, 382 insertions(+), 26 deletions(-) create mode 100644 sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java create mode 100644 sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java index b0854fffc636..52f74c2adcdd 100644 --- a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/ResumableUploadChunkCoordinator.java @@ -39,11 +39,9 @@ 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.IOException; import java.io.InputStream; -import java.util.Arrays; import java.util.concurrent.CancellationException; import org.jspecify.annotations.NullMarked; import org.jspecify.annotations.Nullable; @@ -57,14 +55,10 @@ @NullMarked final class ResumableUploadChunkCoordinator { - private static final byte[] EMPTY_PAYLOAD = new byte[0]; - private final UnaryCallable> 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 uploadResultFuture = SettableApiFuture.create(); private volatile @Nullable ApiFuture inFlightFuture; @@ -78,10 +72,9 @@ final class ResumableUploadChunkCoordinator { 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 getFuture() { @@ -107,35 +100,25 @@ private void transmitChunk(long currentOffset) { } // Read the next chunk slice from the payload stream. - int bytesRead; try { - bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize); + buffer.fill(currentOffset); } catch (IOException e) { uploadResultFuture.setException(e); return; } // 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; + long chunkLength = chunkRequest.getPayload().length; + boolean isFinal = chunkRequest.isFinal(); try { ApiFuture> chunkFuture = uploadChunkCallable.futureCall(chunkRequest, callContext); @@ -191,3 +174,4 @@ private boolean tryRegisterInFlightFuture(ApiFuture future) { return true; } } + diff --git a/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java new file mode 100644 index 000000000000..57d794a0fca1 --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/main/java/com/google/api/gax/rpc/RewindableStreamBuffer.java @@ -0,0 +1,155 @@ +/* + * 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 from the stream starting at {@code targetOffset}, reading up to chunk size. + * + * @param targetOffset the absolute stream offset corresponding to the start of this chunk + * @throws IOException if reading from the stream fails + */ + void fill(long targetOffset) throws IOException { + this.bufferBaseOffset = targetOffset; + 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}. + * + *

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); + } + + int getPayloadLength() { + return payloadLength; + } + + long getBufferBaseOffset() { + return bufferBaseOffset; + } + + boolean isFinal() { + return isFinal; + } +} diff --git a/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java new file mode 100644 index 000000000000..40351740863c --- /dev/null +++ b/sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/RewindableStreamBufferTest.java @@ -0,0 +1,217 @@ +/* + * 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.truth.Truth.assertThat; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.io.ByteArrayInputStream; +import java.io.FilterInputStream; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.StandardCharsets; +import org.junit.jupiter.api.Test; + +class RewindableStreamBufferTest { + + private static final String UPLOAD_URL = "https://upload.example.com/session-1"; + + @Test + void testExactMultiplePayloads() throws IOException { + byte[] data = "0123456789abcdef".getBytes(StandardCharsets.UTF_8); // 16 bytes, chunk size 8 + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + // Chunk 0: 8 bytes + buffer.fill(0L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(buffer.isFinal()).isFalse(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); + + // Chunk 1: 8 bytes + buffer.fill(8L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(8L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(buffer.isFinal()).isFalse(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("89abcdef"); + + // Final 0-byte finalize chunk + buffer.fill(16L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(16L); + assertThat(buffer.getPayloadLength()).isEqualTo(0); + assertThat(buffer.isFinal()).isTrue(); + assertThat(buffer.getPayload()).isEmpty(); + } + + @Test + void testShortFinalChunk() throws IOException { + byte[] data = "short".getBytes(StandardCharsets.UTF_8); // 5 bytes, chunk size 8 + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + buffer.fill(0L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(5); + assertThat(buffer.isFinal()).isTrue(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("short"); + } + + @Test + void testZeroBytePayload() throws IOException { + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(new byte[0]), 8, UPLOAD_URL); + + buffer.fill(0L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(0); + assertThat(buffer.isFinal()).isTrue(); + assertThat(buffer.getPayload()).isEmpty(); + } + + @Test + void testRealignToMidBufferOffset_compactsAndTopsUp() throws IOException { + // 20 bytes: chunk size 8 + byte[] data = "0123456789ABCDEFGHIJ".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + // Initial fill: "01234567" (bytes 0..7) + buffer.fill(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); + + // Server committed 5 bytes (0..4), so next committed offset is 5. + // Remaining uncommitted bytes in buffer: "567" (3 bytes). + // realignTo(5) compacts "567" to buffer[0..3) and tops up 5 more bytes ("89ABC") from stream. + buffer.realignTo(5L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(5L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); // 3 remaining + 5 topped up = 8 + assertThat(buffer.isFinal()).isFalse(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("56789ABC"); + } + + @Test + void testRealignToBufferBaseOffset_isNoOp() throws IOException { + byte[] data = "0123456789".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + buffer.fill(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); + + // Realigning to exactly the buffer base offset (0) is a no-op + buffer.realignTo(0L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); + } + + @Test + void testRealignToBelowBaseOffset_throwsIllegalStateException() throws IOException { + byte[] data = "0123456789abcdef".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + // Advanced to chunk 1 (base offset 8) + buffer.fill(8L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(8L); + + // Server requests offset 4, which is below buffer base offset 8 + IllegalStateException exception = + assertThrows(IllegalStateException.class, () -> buffer.realignTo(4L)); + + assertThat(exception.getMessage()).contains("4"); + assertThat(exception.getMessage()).contains("8"); + assertThat(exception.getMessage()).contains(UPLOAD_URL); + } + + @Test + void testRealignToBeyondBufferWindow_throwsIllegalStateException() throws IOException { + byte[] data = "0123456789abcdef".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + // Initial fill at 0: window is [0, 8] + buffer.fill(0L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + + // Server reports committed offset 10, which is beyond current buffer window [0, 8] + IllegalStateException exception = + assertThrows(IllegalStateException.class, () -> buffer.realignTo(10L)); + + assertThat(exception.getMessage()).contains("10"); + assertThat(exception.getMessage()).contains("8"); + assertThat(exception.getMessage()).contains(UPLOAD_URL); + } + + @Test + void testRealignToMidBufferOffset_reachingEofMarksFinal() throws IOException { + // 10 bytes total, chunk size 8 + byte[] data = "0123456789".getBytes(StandardCharsets.UTF_8); + RewindableStreamBuffer buffer = + new RewindableStreamBuffer(new ByteArrayInputStream(data), 8, UPLOAD_URL); + + // Initial fill: "01234567" (bytes 0..7), stream still has "89" remaining + buffer.fill(0L); + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(buffer.isFinal()).isFalse(); + + // Server committed 4 bytes (0..3). Compacts "4567" (4 bytes) and tops up "89" (2 bytes < space) + buffer.realignTo(4L); + assertThat(buffer.getBufferBaseOffset()).isEqualTo(4L); + assertThat(buffer.getPayloadLength()).isEqualTo(6); + assertThat(buffer.isFinal()).isTrue(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("456789"); + } + + @Test + void testFillWithShortReads_greedilyFillsBufferToCapacity() throws IOException { + byte[] data = "01234567".getBytes(StandardCharsets.UTF_8); // 8 bytes + // Stream that yields at most 2 bytes per read + InputStream shortReadingStream = + new FilterInputStream(new ByteArrayInputStream(data)) { + @Override + public int read(byte[] b, int off, int len) throws IOException { + return super.read(b, off, Math.min(len, 2)); + } + }; + + RewindableStreamBuffer buffer = new RewindableStreamBuffer(shortReadingStream, 8, UPLOAD_URL); + buffer.fill(0L); + + // Must greedily fill all 8 bytes despite short reads, and not be marked final yet + assertThat(buffer.getPayloadLength()).isEqualTo(8); + assertThat(buffer.isFinal()).isFalse(); + assertThat(new String(buffer.getPayload(), StandardCharsets.UTF_8)).isEqualTo("01234567"); + } +}