Skip to content

Commit 074e557

Browse files
whoweswhowes
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 08e807a commit 074e557

7 files changed

Lines changed: 485 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
@@ -46,6 +46,7 @@
4646
import java.io.ByteArrayInputStream;
4747
import java.io.InputStream;
4848
import java.nio.charset.StandardCharsets;
49+
import java.util.Arrays;
4950
import java.util.Collections;
5051
import java.util.List;
5152
import java.util.Map;
@@ -81,7 +82,12 @@ public Map<String, List<String>> getQueryParamNames(ChunkUploadRequest request)
8182

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

8793
@Override
@@ -110,7 +116,7 @@ private ResumableUploadChunkCallable(
110116
public ApiFuture<ChunkUploadResponse<ResponseT>> futureCall(
111117
ChunkUploadRequest request, @Nullable ApiCallContext inputContext) {
112118
Preconditions.checkNotNull(request);
113-
boolean isPayloadEmpty = request.getPayload().length == 0;
119+
boolean isPayloadEmpty = request.getPayloadLength() == 0;
114120
String command;
115121
if (request.isFinal()) {
116122
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: 11 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -39,11 +39,9 @@
3939
import com.google.api.gax.resumable.ChunkUploadRequest;
4040
import com.google.api.gax.resumable.ChunkUploadResponse;
4141
import com.google.api.gax.resumable.ResumableUploadStatus;
42-
import com.google.common.io.ByteStreams;
4342
import com.google.common.util.concurrent.MoreExecutors;
4443
import java.io.IOException;
4544
import java.io.InputStream;
46-
import java.util.Arrays;
4745
import java.util.concurrent.CancellationException;
4846
import java.util.concurrent.Executor;
4947
import org.jspecify.annotations.NullMarked;
@@ -58,18 +56,14 @@
5856
@NullMarked
5957
final class ResumableUploadChunkCoordinator<ResponseT> {
6058

61-
private static final byte[] EMPTY_PAYLOAD = new byte[0];
62-
6359
// Serializes buffer mutations across multiple threads (i.e. from retry/recovery)
6460
private final Executor chunkExecutor =
6561
MoreExecutors.newSequentialExecutor(MoreExecutors.directExecutor());
6662

6763
private final UnaryCallable<ChunkUploadRequest, ChunkUploadResponse<ResponseT>>
6864
uploadChunkCallable;
6965
private final String uploadUrl;
70-
private final InputStream payload;
71-
private final byte[] buffer;
72-
private final int chunkSize;
66+
private final RewindableStreamBuffer buffer;
7367
private final ApiCallContext callContext;
7468
private final SettableApiFuture<ResponseT> result = SettableApiFuture.create();
7569
private volatile @Nullable ApiFuture<?> currentChunkFuture;
@@ -83,10 +77,9 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
8377
this.uploadChunkCallable =
8478
checkNotNull(uploadChunkCallable, "uploadChunkCallable must not be null");
8579
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
86-
this.payload = checkNotNull(payload, "payload must not be null");
87-
this.chunkSize = chunkSize;
80+
checkNotNull(payload, "payload must not be null");
8881
this.callContext = checkNotNull(callContext, "callContext must not be null");
89-
this.buffer = new byte[chunkSize];
82+
this.buffer = new RewindableStreamBuffer(payload, chunkSize, uploadUrl);
9083
}
9184

9285
ApiFuture<ResponseT> start() {
@@ -109,35 +102,26 @@ private void transmitChunk(long currentOffset) {
109102
}
110103

111104
// Read the next chunk slice from the payload stream.
112-
int bytesRead;
113105
try {
114-
bytesRead = ByteStreams.read(payload, buffer, 0, chunkSize);
106+
buffer.fill(currentOffset);
115107
} catch (IOException e) {
116108
result.setException(e);
117109
return;
118110
}
119111

120112
// Determine if this is the final chunk and build the chunk request.
121-
boolean isFinal = bytesRead < chunkSize;
122-
byte[] chunkPayload;
123-
if (bytesRead == chunkSize) {
124-
chunkPayload = buffer;
125-
} else if (bytesRead == 0) {
126-
chunkPayload = EMPTY_PAYLOAD;
127-
} else {
128-
chunkPayload = Arrays.copyOf(buffer, bytesRead);
129-
}
130-
131113
ChunkUploadRequest chunkRequest =
132114
ChunkUploadRequest.newBuilder()
133115
.setUploadUrl(uploadUrl)
134-
.setPayload(chunkPayload)
135-
.setOffset(currentOffset)
136-
.setFinal(isFinal)
116+
.setPayload(buffer.getBuffer())
117+
.setPayloadLength(buffer.getPayloadLength())
118+
.setOffset(buffer.getBufferBaseOffset())
119+
.setFinal(buffer.isFinal())
137120
.build();
138121

139122
// Dispatch the chunk upload call and register the in-flight future for cancellation.
140-
long chunkLength = chunkPayload.length;
123+
long chunkLength = chunkRequest.getPayloadLength();
124+
boolean isFinal = chunkRequest.isFinal();
141125
try {
142126
ApiFuture<ChunkUploadResponse<ResponseT>> chunkFuture =
143127
uploadChunkCallable.futureCall(chunkRequest, callContext);
@@ -183,3 +167,4 @@ public void onFailure(Throwable t) {
183167
}
184168
}
185169
}
170+

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

Lines changed: 164 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,164 @@
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+
import org.jspecify.annotations.Nullable;
40+
41+
/** Manages a single-chunk rewindable buffer over an {@link InputStream} for resumable uploads. */
42+
@NullMarked
43+
final class RewindableStreamBuffer {
44+
45+
private final InputStream inputStream;
46+
private final int chunkSize;
47+
private final String uploadUrl;
48+
private final byte[] buffer;
49+
50+
private long bufferBaseOffset;
51+
private int payloadLength;
52+
private boolean isFinal;
53+
private boolean streamExhausted;
54+
55+
RewindableStreamBuffer(InputStream inputStream, int chunkSize, String uploadUrl) {
56+
this.inputStream = checkNotNull(inputStream, "inputStream must not be null");
57+
checkArgument(chunkSize > 0, "chunkSize must be > 0");
58+
this.chunkSize = chunkSize;
59+
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
60+
this.buffer = new byte[chunkSize];
61+
this.bufferBaseOffset = 0L;
62+
this.payloadLength = 0;
63+
this.isFinal = false;
64+
this.streamExhausted = false;
65+
}
66+
67+
/**
68+
* Advances the buffer from the stream starting at {@code targetOffset}, reading up to chunk size.
69+
*
70+
* @param targetOffset the absolute stream offset corresponding to the start of this chunk
71+
* @throws IOException if reading from the stream fails
72+
*/
73+
void fill(long targetOffset) throws IOException {
74+
this.bufferBaseOffset = targetOffset;
75+
this.payloadLength = ByteStreams.read(inputStream, buffer, 0, chunkSize);
76+
this.isFinal = (payloadLength < chunkSize);
77+
if (this.isFinal) {
78+
this.streamExhausted = true;
79+
}
80+
}
81+
82+
/**
83+
* Realigns the buffer window to {@code committedOffset}.
84+
*
85+
* <p>Compacts forward within the existing buffer to discard already-committed bytes, and then
86+
* tops up the buffer to capacity from the underlying stream.
87+
*
88+
* @param committedOffset the server's committed byte offset
89+
* @throws FailedPreconditionException if {@code committedOffset} is below the buffer's base
90+
* offset or beyond the current buffer window
91+
* @throws IOException if reading from the stream fails
92+
*/
93+
void realignTo(long committedOffset) throws IOException {
94+
if (committedOffset < bufferBaseOffset) {
95+
throw protocolViolation(
96+
String.format(
97+
"Server committed offset %d is below buffer base offset %d for upload URL %s; cannot"
98+
+ " rewind stream before buffer base",
99+
committedOffset, bufferBaseOffset, uploadUrl));
100+
}
101+
102+
if (committedOffset > bufferBaseOffset + payloadLength) {
103+
throw protocolViolation(
104+
String.format(
105+
"Server committed offset %d is beyond current buffer window [%d, %d] for upload URL"
106+
+ " %s",
107+
committedOffset, bufferBaseOffset, bufferBaseOffset + payloadLength, uploadUrl));
108+
}
109+
110+
int committedWithinBuffer = (int) (committedOffset - bufferBaseOffset);
111+
int remainingBytes = payloadLength - committedWithinBuffer;
112+
113+
if (remainingBytes > 0 && committedWithinBuffer > 0) {
114+
System.arraycopy(buffer, committedWithinBuffer, buffer, 0, remainingBytes);
115+
}
116+
117+
this.bufferBaseOffset = committedOffset;
118+
this.payloadLength = remainingBytes;
119+
120+
if (!streamExhausted && payloadLength < chunkSize) {
121+
int space = chunkSize - payloadLength;
122+
int additionalRead = ByteStreams.read(inputStream, buffer, payloadLength, space);
123+
payloadLength += additionalRead;
124+
if (additionalRead < space) {
125+
streamExhausted = true;
126+
}
127+
}
128+
129+
this.isFinal = streamExhausted;
130+
}
131+
132+
byte[] getBuffer() {
133+
return buffer;
134+
}
135+
136+
int getPayloadLength() {
137+
return payloadLength;
138+
}
139+
140+
long getBufferBaseOffset() {
141+
return bufferBaseOffset;
142+
}
143+
144+
boolean isFinal() {
145+
return isFinal;
146+
}
147+
148+
private static FailedPreconditionException protocolViolation(String message) {
149+
return new FailedPreconditionException(message, null, FAILED_PRECONDITION_STATUS_CODE, false);
150+
}
151+
152+
private static final StatusCode FAILED_PRECONDITION_STATUS_CODE =
153+
new StatusCode() {
154+
@Override
155+
public StatusCode.Code getCode() {
156+
return StatusCode.Code.FAILED_PRECONDITION;
157+
}
158+
159+
@Override
160+
public @Nullable Object getTransportCode() {
161+
return null;
162+
}
163+
};
164+
}

0 commit comments

Comments
 (0)