Skip to content

Commit f130ccc

Browse files
committed
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.
1 parent ce8a9f5 commit f130ccc

8 files changed

Lines changed: 661 additions & 33 deletions

File tree

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

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import java.io.ByteArrayInputStream;
4646
import java.io.InputStream;
4747
import java.nio.charset.StandardCharsets;
48+
import java.util.Arrays;
4849
import java.util.Collections;
4950
import java.util.List;
5051
import java.util.Map;
@@ -82,7 +83,12 @@ public Map<String, List<String>> getQueryParamNames(ChunkUploadRequest request)
8283

8384
@Override
8485
public byte[] getBinaryRequestBody(ChunkUploadRequest request) {
85-
return request.getPayload();
86+
int length = request.getPayloadLength();
87+
byte[] payload = request.getPayload();
88+
if (length == payload.length) {
89+
return payload;
90+
}
91+
return Arrays.copyOf(payload, length);
8692
}
8793

8894
@Override
@@ -111,7 +117,7 @@ private ResumableUploadChunkCallable(
111117
public ApiFuture<ChunkUploadResponse<ResponseT>> futureCall(
112118
ChunkUploadRequest request, @Nullable ApiCallContext inputContext) {
113119
Preconditions.checkNotNull(request);
114-
boolean isPayloadEmpty = request.getPayload().length == 0;
120+
boolean isPayloadEmpty = request.getPayloadLength() == 0;
115121
String command;
116122
if (request.isFinal()) {
117123
command = !isPayloadEmpty ? COMMAND_UPLOAD_FINALIZE : COMMAND_FINALIZE;

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

Lines changed: 29 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
import com.google.api.core.BetaApi;
3333
import com.google.api.core.InternalApi;
3434
import com.google.auto.value.AutoValue;
35+
import com.google.common.base.Preconditions;
3536
import org.jspecify.annotations.NullMarked;
3637

3738
/** Request value object for uploading a chunk to an active resumable upload session. */
@@ -48,6 +49,9 @@ public abstract class ChunkUploadRequest {
4849
@SuppressWarnings("mutable")
4950
public abstract byte[] getPayload();
5051

52+
/** The number of bytes within {@link #getPayload()} to upload. */
53+
public abstract int getPayloadLength();
54+
5155
/** The byte offset of this chunk in the overall stream. */
5256
public abstract long getOffset();
5357

@@ -56,8 +60,12 @@ public abstract class ChunkUploadRequest {
5660

5761
public abstract Builder toBuilder();
5862

63+
private static final int UNSET_PAYLOAD_LENGTH = Integer.MIN_VALUE;
64+
5965
public static Builder newBuilder() {
60-
return new AutoValue_ChunkUploadRequest.Builder().setFinal(false);
66+
return new AutoValue_ChunkUploadRequest.Builder()
67+
.setFinal(false)
68+
.setPayloadLength(UNSET_PAYLOAD_LENGTH);
6169
}
6270

6371
@AutoValue.Builder
@@ -66,10 +74,29 @@ public abstract static class Builder {
6674

6775
public abstract Builder setPayload(byte[] payload);
6876

77+
public abstract Builder setPayloadLength(int payloadLength);
78+
6979
public abstract Builder setOffset(long offset);
7080

7181
public abstract Builder setFinal(boolean isFinal);
7282

73-
public abstract ChunkUploadRequest build();
83+
abstract byte[] getPayload();
84+
85+
abstract int getPayloadLength();
86+
87+
abstract ChunkUploadRequest autoBuild();
88+
89+
public ChunkUploadRequest build() {
90+
if (getPayloadLength() == UNSET_PAYLOAD_LENGTH) {
91+
setPayloadLength(getPayload().length);
92+
}
93+
ChunkUploadRequest request = autoBuild();
94+
Preconditions.checkArgument(
95+
request.getPayloadLength() >= 0, "payloadLength must be non-negative");
96+
Preconditions.checkArgument(
97+
request.getPayloadLength() <= request.getPayload().length,
98+
"payloadLength exceeds payload array length");
99+
return request;
100+
}
74101
}
75102
}

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

Lines changed: 19 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -44,13 +44,11 @@
4444
import com.google.api.gax.retrying.RetrySettings;
4545
import com.google.api.gax.retrying.RetryingFuture;
4646
import com.google.api.gax.retrying.ScheduledRetryingExecutor;
47-
import com.google.common.io.ByteStreams;
4847
import com.google.common.util.concurrent.MoreExecutors;
4948
import com.google.errorprone.annotations.concurrent.GuardedBy;
5049
import java.io.IOException;
5150
import java.io.InputStream;
5251
import java.time.Duration;
53-
import java.util.Arrays;
5452
import java.util.concurrent.CancellationException;
5553
import java.util.concurrent.atomic.AtomicBoolean;
5654
import java.util.concurrent.atomic.AtomicLong;
@@ -77,8 +75,6 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
7775
.setMaxAttempts(5)
7876
.build();
7977

80-
private static final byte[] EMPTY_PAYLOAD = new byte[0];
81-
8278
private final Object lock = new Object();
8379
private final AtomicBoolean dispatching = new AtomicBoolean(false);
8480
private final AtomicLong nextChunkOffset = new AtomicLong(-1L);
@@ -88,13 +84,13 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
8884
private final RetryingCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
8985
retryingChunkCallable;
9086
private final InputStream payload;
91-
private final byte[] buffer;
9287
private final int chunkSize;
9388
private final ApiCallContext callContext;
9489
private final ClientContext clientContext;
9590
private final RetrySettings chunkRetrySettings;
9691

9792
private volatile @Nullable String uploadSessionUrl;
93+
private volatile @Nullable RewindableStreamBuffer buffer;
9894

9995
@GuardedBy("lock")
10096
private boolean done;
@@ -123,7 +119,6 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
123119
this.callContext = checkNotNull(callContext, "callContext must not be null");
124120
this.clientContext = checkNotNull(clientContext, "clientContext must not be null");
125121
this.chunkRetrySettings = DEFAULT_CHUNK_RETRY_SETTINGS;
126-
this.buffer = new byte[chunkSize];
127122

128123
RetryAlgorithm<ChunkUploadResponse<ResponseT>> retryAlgorithm =
129124
new RetryAlgorithm<>(
@@ -152,6 +147,7 @@ public void onSuccess(ResumableUploadSession session) {
152147
}
153148
}
154149
uploadSessionUrl = session.getUploadUrl();
150+
buffer = new RewindableStreamBuffer(payload, chunkSize, uploadSessionUrl);
155151
scheduleNextChunk(0L);
156152
}
157153

@@ -258,43 +254,40 @@ private void transmitSingleChunk(long currentOffset) {
258254
}
259255
}
260256

261-
int bytesRead;
262-
try {
263-
bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize);
264-
} catch (IOException e) {
265-
finish(null, e);
257+
String url = uploadSessionUrl;
258+
if (url == null) {
259+
finish(null, new IllegalStateException("Upload session URL not available"));
266260
return;
267261
}
268262

269-
boolean isFinal = bytesRead < chunkSize;
270-
byte[] chunkPayload;
271-
if (bytesRead == chunkSize) {
272-
chunkPayload = buffer;
273-
} else if (bytesRead == 0) {
274-
chunkPayload = EMPTY_PAYLOAD;
275-
} else {
276-
chunkPayload = Arrays.copyOf(buffer, bytesRead);
263+
RewindableStreamBuffer streamBuffer = buffer;
264+
if (streamBuffer == null) {
265+
finish(null, new IllegalStateException("Upload buffer not initialized"));
266+
return;
277267
}
278268

279-
String url = uploadSessionUrl;
280-
if (url == null) {
281-
finish(null, new IllegalStateException("Upload session URL not available"));
269+
try {
270+
streamBuffer.fill(currentOffset);
271+
} catch (IOException e) {
272+
finish(null, e);
282273
return;
283274
}
284275

285276
ChunkUploadRequest chunkRequest =
286277
ChunkUploadRequest.newBuilder()
287278
.setUploadUrl(url)
288-
.setPayload(chunkPayload)
289-
.setOffset(currentOffset)
290-
.setFinal(isFinal)
279+
.setPayload(streamBuffer.getBuffer())
280+
.setPayloadLength(streamBuffer.getPayloadLength())
281+
.setOffset(streamBuffer.getBufferBaseOffset())
282+
.setFinal(streamBuffer.isFinal())
291283
.build();
292284

293285
RetryingFuture<ChunkUploadResponse<ResponseT>> retryingFuture =
294286
retryingChunkCallable.futureCall(chunkRequest, callContext);
295287
setInFlightFuture(retryingFuture);
296288

297-
long chunkLength = chunkPayload.length;
289+
long chunkLength = chunkRequest.getPayloadLength();
290+
boolean isFinal = chunkRequest.isFinal();
298291
ApiFutures.addCallback(
299292
retryingFuture,
300293
new ApiFutureCallback<ChunkUploadResponse<ResponseT>>() {
Lines changed: 157 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,157 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Redistribution and use in source and binary forms, with or without
5+
* modification, are permitted provided that the following conditions are
6+
* met:
7+
*
8+
* * Redistributions of source code must retain the above copyright
9+
* notice, this list of conditions and the following disclaimer.
10+
* * Redistributions in binary form must reproduce the above
11+
* copyright notice, this list of conditions and the following disclaimer
12+
* in the documentation and/or other materials provided with the
13+
* distribution.
14+
* * Neither the name of Google LLC nor the names of its
15+
* contributors may be used to endorse or promote products derived from
16+
* this software without specific prior written permission.
17+
*
18+
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
19+
* "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
20+
* LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
21+
* A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
22+
* OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
23+
* SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
24+
* LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
25+
* DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
26+
* THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
27+
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28+
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
29+
*/
30+
package com.google.api.gax.rpc;
31+
32+
import static com.google.common.base.Preconditions.checkArgument;
33+
import static com.google.common.base.Preconditions.checkNotNull;
34+
35+
import com.google.common.io.ByteStreams;
36+
import java.io.IOException;
37+
import java.io.InputStream;
38+
import org.jspecify.annotations.NullMarked;
39+
40+
/**
41+
* Manages a single-chunk buffer over an {@link InputStream} for resumable uploads.
42+
*
43+
* <p>The buffer holds at most one chunk of data in a reused backing array. It supports forward
44+
* compaction and topping up upon recovery realignment, and enforces the boundary condition that
45+
* requests to rewind before the buffer's base offset fail with an unrecoverable {@link
46+
* FailedPreconditionException}.
47+
*/
48+
@NullMarked
49+
final class RewindableStreamBuffer {
50+
51+
private final InputStream inputStream;
52+
private final int chunkSize;
53+
private final String uploadUrl;
54+
private final byte[] buffer;
55+
56+
private long bufferBaseOffset;
57+
private int payloadLength;
58+
private boolean isFinal;
59+
private boolean streamExhausted;
60+
61+
RewindableStreamBuffer(InputStream inputStream, int chunkSize, String uploadUrl) {
62+
this.inputStream = checkNotNull(inputStream, "inputStream must not be null");
63+
checkArgument(chunkSize > 0, "chunkSize must be > 0");
64+
this.chunkSize = chunkSize;
65+
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
66+
this.buffer = new byte[chunkSize];
67+
this.bufferBaseOffset = 0L;
68+
this.payloadLength = 0;
69+
this.isFinal = false;
70+
this.streamExhausted = false;
71+
}
72+
73+
/**
74+
* Advances the buffer from the stream starting at {@code targetOffset}, reading up to chunk size.
75+
*
76+
* @param targetOffset the absolute stream offset corresponding to the start of this chunk
77+
* @throws IOException if reading from the stream fails
78+
*/
79+
void fill(long targetOffset) throws IOException {
80+
this.bufferBaseOffset = targetOffset;
81+
this.payloadLength = ByteStreams.read(inputStream, buffer, 0, chunkSize);
82+
this.isFinal = (payloadLength < chunkSize);
83+
if (this.isFinal) {
84+
this.streamExhausted = true;
85+
}
86+
}
87+
88+
/**
89+
* Realigns the buffer window to {@code committedOffset}.
90+
*
91+
* <p>Compacts forward within the existing buffer to discard already-committed bytes, and then
92+
* tops up the buffer to capacity from the underlying stream.
93+
*
94+
* @param committedOffset the server's committed byte offset
95+
* @throws FailedPreconditionException if {@code committedOffset} is below the buffer's base
96+
* offset or beyond the current buffer window
97+
* @throws IOException if reading from the stream fails
98+
*/
99+
void realignTo(long committedOffset) throws IOException {
100+
if (committedOffset < bufferBaseOffset) {
101+
throw UploadErrors.protocolViolation(
102+
String.format(
103+
"Server committed offset %d is below buffer base offset %d for upload URL %s; cannot"
104+
+ " rewind stream before buffer base",
105+
committedOffset, bufferBaseOffset, uploadUrl));
106+
}
107+
108+
if (committedOffset > bufferBaseOffset + payloadLength) {
109+
throw UploadErrors.protocolViolation(
110+
String.format(
111+
"Server committed offset %d is beyond current buffer window [%d, %d] for upload URL"
112+
+ " %s",
113+
committedOffset, bufferBaseOffset, bufferBaseOffset + payloadLength, uploadUrl));
114+
}
115+
116+
int committedWithinBuffer = (int) (committedOffset - bufferBaseOffset);
117+
int remainingBytes = payloadLength - committedWithinBuffer;
118+
119+
if (remainingBytes > 0 && committedWithinBuffer > 0) {
120+
System.arraycopy(buffer, committedWithinBuffer, buffer, 0, remainingBytes);
121+
}
122+
123+
this.bufferBaseOffset = committedOffset;
124+
this.payloadLength = remainingBytes;
125+
126+
if (!streamExhausted && payloadLength < chunkSize) {
127+
int space = chunkSize - payloadLength;
128+
int additionalRead = ByteStreams.read(inputStream, buffer, payloadLength, space);
129+
payloadLength += additionalRead;
130+
if (additionalRead < space) {
131+
streamExhausted = true;
132+
}
133+
}
134+
135+
this.isFinal = streamExhausted;
136+
}
137+
138+
byte[] getBuffer() {
139+
return buffer;
140+
}
141+
142+
int getPayloadLength() {
143+
return payloadLength;
144+
}
145+
146+
long getBufferBaseOffset() {
147+
return bufferBaseOffset;
148+
}
149+
150+
boolean isFinal() {
151+
return isFinal;
152+
}
153+
154+
boolean isEmpty() {
155+
return payloadLength == 0;
156+
}
157+
}

0 commit comments

Comments
 (0)