Skip to content

Commit 0ad413f

Browse files
whoweswhowes
committed
feat(gax): align resumable upload chunk size to server granularity and surface effective chunk size
1 parent f1db2bb commit 0ad413f

8 files changed

Lines changed: 129 additions & 25 deletions

File tree

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
6868
private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
6969
queryStatusCallable;
7070
private final String uploadUrl;
71+
private final int chunkSize;
7172
private final RewindableStreamBuffer buffer;
7273
private final ApiCallContext callContext;
7374
private final ExponentialRetryAlgorithm recoveryAlgorithm;
@@ -93,6 +94,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
9394
this.queryStatusCallable =
9495
checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
9596
this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
97+
this.chunkSize = chunkSize;
9698
checkNotNull(payload, "payload must not be null");
9799
this.callContext = checkNotNull(callContext, "callContext must not be null");
98100
this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null");
@@ -107,7 +109,7 @@ ApiFuture<ResponseT> getFuture() {
107109
}
108110

109111
void start() {
110-
progressTracker.onStarted(uploadUrl);
112+
progressTracker.onStarted(uploadUrl, chunkSize);
111113
uploadResultFuture.addListener(
112114
() -> {
113115
Future<?> inFlight = inFlightFuture;

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131

3232
import com.google.api.core.ApiFuture;
3333
import com.google.api.core.BetaApi;
34+
import com.google.api.core.InternalExtensionOnly;
3435
import org.jspecify.annotations.NullMarked;
3536
import org.jspecify.annotations.Nullable;
3637

@@ -43,9 +44,16 @@
4344
* @param <ResponseT> the type of the final response message returned once the upload completes
4445
*/
4546
@BetaApi
47+
@InternalExtensionOnly
4648
@NullMarked
4749
public interface ResumableUploadFuture<ResponseT> extends ApiFuture<ResponseT> {
4850

4951
/** Returns the upload session URL, or {@code null} if session initiation is in progress. */
5052
@Nullable String getUploadSessionUrl();
53+
54+
/**
55+
* Returns the effective chunk size in bytes aligned to the server's chunk granularity, or {@code
56+
* null} if session initiation is in progress.
57+
*/
58+
@Nullable Integer getEffectiveChunkSize();
5159
}

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

Lines changed: 31 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
8686

8787
private volatile @Nullable InputStream payload;
8888
private volatile @Nullable String uploadSessionUrl;
89+
private volatile @Nullable Integer effectiveChunkSize;
8990

9091
// Tracks the current operation's Future (start, chunk upload) to propagate cancellation, or
9192
// null once a terminal transition (succeed, fail, cancel) has claimed completion.
@@ -161,13 +162,20 @@ private void start(InputStreamSupplier payloadSupplier) throws IOException {
161162
@Override
162163
public void onSuccess(ResumableUploadSession session) {
163164
String sessionUrl = session.getUploadUrl();
165+
int chunkSize;
166+
try {
167+
chunkSize = alignChunkSize(options.getChunkSize(), session.getChunkGranularity());
168+
} catch (Throwable t) {
169+
fail(t);
170+
return;
171+
}
164172
ResumableUploadChunkCoordinator<ResponseT> coordinator =
165173
new ResumableUploadChunkCoordinator<>(
166174
uploadChunkCallable,
167175
queryStatusCallable,
168176
sessionUrl,
169177
payload,
170-
options.getChunkSize(),
178+
chunkSize,
171179
callContext,
172180
recoveryAlgorithm,
173181
executor,
@@ -178,6 +186,7 @@ public void onSuccess(ResumableUploadSession session) {
178186
return;
179187
}
180188
uploadSessionUrl = sessionUrl;
189+
effectiveChunkSize = chunkSize;
181190
inFlightFuture = uploadFuture;
182191
}
183192
ApiFutures.addCallback(
@@ -256,9 +265,12 @@ private ApiException toTerminalException(Throwable t) {
256265
}
257266
String message = firstNonNull(t.getMessage(), t.getClass().getSimpleName());
258267
if (uploadSessionUrl != null) {
259-
message = message + " (upload URL: " + uploadSessionUrl + ")";
268+
message =
269+
String.format(
270+
"%s (upload URL: %s, effective chunk size: %d)",
271+
message, uploadSessionUrl, effectiveChunkSize);
260272
}
261-
// Rethrow with URL included in message.
273+
// Rethrow with URL and effective chunk size included in message.
262274
if (t instanceof ApiException) {
263275
ApiException apiException = (ApiException) t;
264276
return ApiExceptionFactory.createException(
@@ -272,6 +284,17 @@ private ApiException toTerminalException(Throwable t) {
272284
message, t, statusCodeOf(StatusCode.Code.UNKNOWN), false);
273285
}
274286

287+
static int alignChunkSize(int configuredChunkSize, long chunkGranularity) {
288+
checkArgument(configuredChunkSize > 0, "configuredChunkSize must be > 0");
289+
checkArgument(chunkGranularity > 0, "chunkGranularity must be > 0");
290+
int granularity = Math.toIntExact(chunkGranularity);
291+
if (configuredChunkSize < granularity) {
292+
return granularity;
293+
}
294+
int unalignedRemainder = configuredChunkSize % granularity;
295+
return configuredChunkSize - unalignedRemainder;
296+
}
297+
275298
private void closePayload(@Nullable Throwable primaryException) {
276299
if (payload == null) {
277300
return;
@@ -290,6 +313,11 @@ private void closePayload(@Nullable Throwable primaryException) {
290313
return uploadSessionUrl;
291314
}
292315

316+
@Override
317+
public @Nullable Integer getEffectiveChunkSize() {
318+
return effectiveChunkSize;
319+
}
320+
293321
@Override
294322
public void addListener(Runnable listener, Executor executor) {
295323
resultFuture.addListener(listener, executor);

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -97,8 +97,9 @@ public static Builder newBuilder() {
9797
public abstract static class Builder {
9898
/**
9999
* Sets the chunk size in bytes. The payload is sent in chunks of this size, and each chunk is
100-
* buffered in memory so it can be resent if the upload needs to recover. Must be positive;
101-
* defaults to 8 MB.
100+
* buffered in memory so it can be resent if the upload needs to recover. If the server
101+
* specifies a chunk granularity, the effective chunk size is aligned to a multiple of that
102+
* granularity. Must be positive; defaults to 8 MB.
102103
*/
103104
public abstract Builder setChunkSize(int chunkSize);
104105

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,12 @@ public abstract class ResumableUploadProgress {
5555
*/
5656
public abstract @Nullable String getUploadUrl();
5757

58+
/**
59+
* Returns the effective chunk size in bytes aligned to the server's chunk granularity, or {@code
60+
* null} if session initiation is pending.
61+
*/
62+
public abstract @Nullable Integer getEffectiveChunkSize();
63+
5864
/** Returns the number of bytes confirmed as uploaded to the server so far. */
5965
public abstract long getBytesUploaded();
6066

@@ -87,6 +93,8 @@ static Builder newBuilder() {
8793
abstract static class Builder {
8894
abstract Builder setUploadUrl(@Nullable String uploadUrl);
8995

96+
abstract Builder setEffectiveChunkSize(@Nullable Integer effectiveChunkSize);
97+
9098
abstract Builder setBytesUploaded(long bytesUploaded);
9199

92100
abstract Builder setState(String state);

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

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,10 +116,14 @@ ResumableUploadProgress getProgress() {
116116
}
117117
}
118118

119-
void onStarted(String uploadUrl) {
119+
void onStarted(String uploadUrl, int effectiveChunkSize) {
120120
checkNotNull(uploadUrl, "uploadUrl must not be null");
121121
transition(
122-
builder -> builder.setState(ResumableUploadProgress.STATE_STARTED).setUploadUrl(uploadUrl));
122+
builder ->
123+
builder
124+
.setState(ResumableUploadProgress.STATE_STARTED)
125+
.setUploadUrl(uploadUrl)
126+
.setEffectiveChunkSize(effectiveChunkSize));
123127
}
124128

125129
void onChunkUploaded(long bytesUploaded) {

‎sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadCallableImplTest.java‎

Lines changed: 63 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1141,20 +1141,22 @@ void testProgressListener_prescribedStateTransitions() throws Exception {
11411141
ResumableUploadFuture<String> future =
11421142
callable.futureCall("resource-path", payloadOf("01234567890123456789"), options);
11431143

1144+
assertThat(future.getEffectiveChunkSize()).isNull();
1145+
11441146
String url = "https://upload.url/progress-transitions";
11451147
startFuture.set(ResumableUploadSession.newBuilder().setUploadUrl(url).build());
11461148

11471149
assertThat(future.get()).isEqualTo("done");
11481150

11491151
assertThat(listener.awaitTerminal())
11501152
.containsExactly(
1151-
createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null),
1152-
createProgress(ResumableUploadProgress.STATE_STARTED, 0L, url),
1153-
createProgress(ResumableUploadProgress.STATE_UPLOADING, 8L, url),
1154-
createProgress(ResumableUploadProgress.STATE_RECOVERING, 8L, url),
1155-
createProgress(ResumableUploadProgress.STATE_OFFSET_RECEIVED, 8L, url),
1156-
createProgress(ResumableUploadProgress.STATE_UPLOADING, 16L, url),
1157-
createProgress(ResumableUploadProgress.STATE_FINALIZED, 20L, url))
1153+
createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null, null),
1154+
createProgress(ResumableUploadProgress.STATE_STARTED, 0L, url, 8),
1155+
createProgress(ResumableUploadProgress.STATE_UPLOADING, 8L, url, 8),
1156+
createProgress(ResumableUploadProgress.STATE_RECOVERING, 8L, url, 8),
1157+
createProgress(ResumableUploadProgress.STATE_OFFSET_RECEIVED, 8L, url, 8),
1158+
createProgress(ResumableUploadProgress.STATE_UPLOADING, 16L, url, 8),
1159+
createProgress(ResumableUploadProgress.STATE_FINALIZED, 20L, url, 8))
11581160
.inOrder();
11591161
}
11601162

@@ -1175,8 +1177,8 @@ void testProgressListener_uploadFailure_transitionsToFailed() throws Exception {
11751177

11761178
assertThat(listener.awaitTerminal())
11771179
.containsExactly(
1178-
createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null),
1179-
createProgress(ResumableUploadProgress.STATE_FAILED, 0L, null))
1180+
createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null, null),
1181+
createProgress(ResumableUploadProgress.STATE_FAILED, 0L, null, null))
11801182
.inOrder();
11811183
}
11821184

@@ -1199,7 +1201,7 @@ void testProgressListener_cancel_transitionsToFailed() throws Exception {
11991201
assertThat(receivedStatuses.get(receivedStatuses.size() - 1))
12001202
.isEqualTo(
12011203
createProgress(
1202-
ResumableUploadProgress.STATE_FAILED, 0L, "https://upload.url/cancel-progress"));
1204+
ResumableUploadProgress.STATE_FAILED, 0L, "https://upload.url/cancel-progress", 8));
12031205
}
12041206

12051207
@Test
@@ -1228,15 +1230,20 @@ void testProgressListener_queryFinal_transitionsToFinalized() throws Exception {
12281230
createProgress(
12291231
ResumableUploadProgress.STATE_FINALIZED,
12301232
5L,
1231-
"https://upload.url/query-final-progress"));
1233+
"https://upload.url/query-final-progress",
1234+
8));
12321235
}
12331236

12341237
private static ResumableUploadProgress createProgress(
1235-
String state, long bytesUploaded, @Nullable String uploadUrl) {
1238+
String state,
1239+
long bytesUploaded,
1240+
@Nullable String uploadUrl,
1241+
@Nullable Integer effectiveChunkSize) {
12361242
return ResumableUploadProgress.newBuilder()
12371243
.setState(state)
12381244
.setBytesUploaded(bytesUploaded)
12391245
.setUploadUrl(uploadUrl)
1246+
.setEffectiveChunkSize(effectiveChunkSize)
12401247
.build();
12411248
}
12421249

@@ -1273,6 +1280,7 @@ void testActionableErrors_startFailure_preservesOriginalException() {
12731280
ExecutionException ex = assertThrows(ExecutionException.class, future::get);
12741281
assertThat(ex.getCause()).isSameInstanceAs(startError);
12751282
assertThat(future.getUploadSessionUrl()).isNull();
1283+
assertThat(future.getEffectiveChunkSize()).isNull();
12761284
}
12771285

12781286
@Test
@@ -1298,9 +1306,52 @@ void testActionableErrors_preservesErrorDetailsAndCauseChain() {
12981306
assertThat(ex.getCause()).isInstanceOf(ApiException.class);
12991307
ApiException cause = (ApiException) ex.getCause();
13001308
assertThat(cause.getMessage()).contains(sessionUrl);
1309+
assertThat(cause.getMessage()).contains("effective chunk size: 8");
13011310
assertThat(cause.getCause()).isSameInstanceAs(transportCause);
13021311
assertThat(cause.getErrorDetails()).isSameInstanceAs(errorDetails);
13031312
assertThat(future.getUploadSessionUrl()).isEqualTo(sessionUrl);
1313+
assertThat(future.getEffectiveChunkSize()).isEqualTo(8);
1314+
}
1315+
1316+
@Test
1317+
void testAlignChunkSize() {
1318+
assertThat(ResumableUploadFutureImpl.alignChunkSize(8, 1L)).isEqualTo(8);
1319+
assertThat(ResumableUploadFutureImpl.alignChunkSize(8, 4L)).isEqualTo(8);
1320+
assertThat(ResumableUploadFutureImpl.alignChunkSize(10, 4L)).isEqualTo(8);
1321+
assertThat(ResumableUploadFutureImpl.alignChunkSize(4, 6L)).isEqualTo(6);
1322+
}
1323+
1324+
@Test
1325+
void testChunkGranularity_alignsChunkSizeAndSurfacesEffectiveSize() throws Exception {
1326+
when(mockStartCallable.futureCall(any(), any()))
1327+
.thenReturn(
1328+
ApiFutures.immediateFuture(
1329+
ResumableUploadSession.newBuilder()
1330+
.setUploadUrl("https://upload.url/granularity-align")
1331+
.setChunkGranularity(4L)
1332+
.build()));
1333+
when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any()))
1334+
.thenReturn(
1335+
ApiFutures.immediateFuture(
1336+
ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null)))
1337+
.thenReturn(
1338+
ApiFutures.immediateFuture(
1339+
ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "done")));
1340+
1341+
// Configured chunkSize = 10, server granularity = 4 -> aligned down to 8
1342+
ResumableUploadOptions options = defaultSettings.toBuilder().setChunkSize(10).build();
1343+
1344+
ResumableUploadFuture<String> future =
1345+
callable.futureCall("resource-path", payloadOf("0123456789ab"), options);
1346+
1347+
assertThat(future.get()).isEqualTo("done");
1348+
assertThat(future.getEffectiveChunkSize()).isEqualTo(8);
1349+
1350+
ArgumentCaptor<ChunkUploadRequest> captor = ArgumentCaptor.forClass(ChunkUploadRequest.class);
1351+
verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any());
1352+
List<ChunkUploadRequest> requests = captor.getAllValues();
1353+
assertChunk(requests.get(0), 0, 8, false);
1354+
assertChunk(requests.get(1), 8, 4, true);
13041355
}
13051356

13061357
@Test

‎sdk-platform-java/gax-java/gax/src/test/java/com/google/api/gax/rpc/ResumableUploadProgressTrackerTest.java‎

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -55,12 +55,13 @@ void testInitialSnapshotOnSubscribe() {
5555
assertThat(initial.getState()).isEqualTo(ResumableUploadProgress.STATE_STARTING);
5656
assertThat(initial.getBytesUploaded()).isEqualTo(0L);
5757
assertThat(initial.getUploadUrl()).isNull();
58+
assertThat(initial.getEffectiveChunkSize()).isNull();
5859
}
5960

6061
@Test
6162
void testProgressUpdates() {
6263
ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
63-
tracker.onStarted("https://upload.url/test");
64+
tracker.onStarted("https://upload.url/test", 256);
6465

6566
List<Long> byteUpdates = new ArrayList<>();
6667
tracker.addListener(
@@ -73,12 +74,13 @@ void testProgressUpdates() {
7374
// Advance to 200 bytes
7475
tracker.onChunkUploaded(200L);
7576
assertThat(byteUpdates).containsExactly(0L, 100L, 200L).inOrder();
77+
assertThat(tracker.getProgress().getEffectiveChunkSize()).isEqualTo(256);
7678
}
7779

7880
@Test
7981
void testListenerExceptionSafety_doesNotDisruptSubsequentUpdates() {
8082
ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
81-
tracker.onStarted("https://upload.url/test");
83+
tracker.onStarted("https://upload.url/test", 256);
8284

8385
AtomicInteger errorCount = new AtomicInteger(0);
8486
List<String> safeReceived = new ArrayList<>();
@@ -110,7 +112,7 @@ void testListenerExceptionSafety_doesNotDisruptSubsequentUpdates() {
110112
@Test
111113
void testTerminalState_noNotificationsAfterFinalized() {
112114
ResumableUploadProgressTracker tracker = new ResumableUploadProgressTracker();
113-
tracker.onStarted("https://upload.url/test");
115+
tracker.onStarted("https://upload.url/test", 512);
114116

115117
List<String> received = new ArrayList<>();
116118
tracker.addListener(status -> received.add(status.getState()), MoreExecutors.directExecutor());
@@ -139,7 +141,7 @@ void testListenerReentrancy_canQueryStatusInsideCallback() {
139141
tracker.addListener(
140142
status -> observed.set(tracker.getProgress()), MoreExecutors.directExecutor());
141143

142-
tracker.onStarted("https://upload.url/reentrancy");
144+
tracker.onStarted("https://upload.url/reentrancy", 1024);
143145
assertThat(observed.get()).isNotNull();
144146
assertThat(observed.get().getState()).isEqualTo(ResumableUploadProgress.STATE_STARTED);
145147
assertThat(observed.get().getUploadUrl()).isEqualTo("https://upload.url/reentrancy");
@@ -162,7 +164,7 @@ void testCustomExecutorDispatch() throws Exception {
162164
},
163165
executor);
164166

165-
tracker.onStarted("https://upload.url/custom-executor");
167+
tracker.onStarted("https://upload.url/custom-executor", 256);
166168
assertThat(latch.await(5, TimeUnit.SECONDS)).isTrue();
167169

168170
synchronized (states) {

0 commit comments

Comments
 (0)