feat(gax): align resumable upload chunk size to server granularity and surface effective chunk size

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

@@ -68,6 +68,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
   private final UnaryCallable<QueryStatusRequest, QueryStatusResponse<ResponseT>>
       queryStatusCallable;
   private final String uploadUrl;
+  private final int chunkSize;
   private final RewindableStreamBuffer buffer;
   private final ApiCallContext callContext;
   private final ExponentialRetryAlgorithm recoveryAlgorithm;
@@ -93,6 +94,7 @@ final class ResumableUploadChunkCoordinator<ResponseT> {
     this.queryStatusCallable =
         checkNotNull(queryStatusCallable, "queryStatusCallable must not be null");
     this.uploadUrl = checkNotNull(uploadUrl, "uploadUrl must not be null");
+    this.chunkSize = chunkSize;
     checkNotNull(payload, "payload must not be null");
     this.callContext = checkNotNull(callContext, "callContext must not be null");
     this.recoveryAlgorithm = checkNotNull(recoveryAlgorithm, "recoveryAlgorithm must not be null");
@@ -107,7 +109,7 @@ ApiFuture<ResponseT> getFuture() {
   }
 
   void start() {
-    progressTracker.onStarted(uploadUrl);
+    progressTracker.onStarted(uploadUrl, chunkSize);
     uploadResultFuture.addListener(
         () -> {
           Future<?> inFlight = inFlightFuture;

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

@@ -31,6 +31,7 @@
 
 import com.google.api.core.ApiFuture;
 import com.google.api.core.BetaApi;
+import com.google.api.core.InternalExtensionOnly;
 import org.jspecify.annotations.NullMarked;
 import org.jspecify.annotations.Nullable;
 
@@ -43,9 +44,16 @@
  * @param <ResponseT> the type of the final response message returned once the upload completes
  */
 @BetaApi
+@InternalExtensionOnly
 @NullMarked
 public interface ResumableUploadFuture<ResponseT> extends ApiFuture<ResponseT> {
 
   /** Returns the upload session URL, or {@code null} if session initiation is in progress. */
   @Nullable String getUploadSessionUrl();
+
+  /**
+   * Returns the effective chunk size in bytes aligned to the server's chunk granularity, or {@code
+   * null} if session initiation is in progress.
+   */
+  @Nullable Integer getEffectiveChunkSize();
 }

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

@@ -85,7 +85,12 @@
   private final SettableApiFuture<ResponseT> resultFuture = SettableApiFuture.create();
 
   private volatile @Nullable InputStream payload;
-  private volatile @Nullable String uploadSessionUrl;
+
+  @GuardedBy("lock")
+  private @Nullable String uploadSessionUrl;
+
+  @GuardedBy("lock")
+  private @Nullable Integer effectiveChunkSize;
 
   // Tracks the current operation's Future (start, chunk upload) to propagate cancellation, or
   // null once a terminal transition (succeed, fail, cancel) has claimed completion.
@@ -161,13 +166,20 @@
           @Override
           public void onSuccess(ResumableUploadSession session) {
             String sessionUrl = session.getUploadUrl();
+            int chunkSize;
+            try {
+              chunkSize = alignChunkSize(options.getChunkSize(), session.getChunkGranularity());
+            } catch (Throwable t) {
+              fail(t);
+              return;
+            }
             ResumableUploadChunkCoordinator<ResponseT> coordinator =
                 new ResumableUploadChunkCoordinator<>(
                     uploadChunkCallable,
                     queryStatusCallable,
                     sessionUrl,
                     payload,
-                    options.getChunkSize(),
+                    chunkSize,
                     callContext,
                     recoveryAlgorithm,
                     executor,
@@ -177,6 +189,7 @@
               if (inFlightFuture == null) {
                 return;
               }
+              effectiveChunkSize = chunkSize;
               uploadSessionUrl = sessionUrl;
               inFlightFuture = uploadFuture;
             }
@@ -213,7 +226,7 @@
 
   private void onTimeout() {
     String message;
-    if (uploadSessionUrl != null) {
+    if (getUploadSessionUrl() != null) {
       message = "Resumable upload timed out";
     } else {
       message = "Resumable upload timed out before session initiation completed";
@@ -251,14 +264,22 @@
   }
 
   private ApiException toTerminalException(Throwable t) {
-    if (t instanceof ApiException && uploadSessionUrl == null) {
+    String sessionUrl;
+    Integer chunkSize;
+    synchronized (lock) {
+      sessionUrl = uploadSessionUrl;
+      chunkSize = effectiveChunkSize;
+    }
+    if (t instanceof ApiException && sessionUrl == null) {
       return (ApiException) t;
     }
     String message = firstNonNull(t.getMessage(), t.getClass().getSimpleName());
-    if (uploadSessionUrl != null) {
-      message = message + " (upload URL: " + uploadSessionUrl + ")";
+    if (sessionUrl != null) {
+      message =
+          String.format(
+              "%s (upload URL: %s, effective chunk size: %d)", message, sessionUrl, chunkSize);
     }
-    // Rethrow with URL included in message.
+    // Rethrow with URL and effective chunk size included in message.
     if (t instanceof ApiException) {
       ApiException apiException = (ApiException) t;
       return ApiExceptionFactory.createException(
@@ -272,6 +293,18 @@
         message, t, statusCodeOf(StatusCode.Code.UNKNOWN), false);
   }
 
+  static int alignChunkSize(int configuredChunkSize, long chunkGranularity) {
+    checkArgument(configuredChunkSize > 0, "configuredChunkSize must be > 0");
+    checkArgument(chunkGranularity > 0, "chunkGranularity must be > 0");
+    checkArgument(
+        configuredChunkSize >= chunkGranularity,
+        "Configured chunk size (%s bytes) must be at least the server chunk granularity (%s bytes)",
+        configuredChunkSize,
+        chunkGranularity);
+    int unalignedRemainder = (int) (configuredChunkSize % chunkGranularity);
+    return configuredChunkSize - unalignedRemainder;
+  }
+
   private void closePayload(@Nullable Throwable primaryException) {
     if (payload == null) {
       return;
@@ -287,7 +320,16 @@
 
   @Override
   public @Nullable String getUploadSessionUrl() {
-    return uploadSessionUrl;
+    synchronized (lock) {
+      return uploadSessionUrl;
+    }
+  }
+
+  @Override
+  public @Nullable Integer getEffectiveChunkSize() {
+    synchronized (lock) {
+      return effectiveChunkSize;
+    }
   }
 
   @Override

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

@@ -97,8 +97,9 @@ public static Builder newBuilder() {
   public abstract static class Builder {
     /**
      * Sets the chunk size in bytes. The payload is sent in chunks of this size, and each chunk is
-     * buffered in memory so it can be resent if the upload needs to recover. Must be positive;
-     * defaults to 8 MB.
+     * buffered in memory so it can be resent if the upload needs to recover. If the server
+     * specifies a chunk granularity, the effective chunk size is rounded down to a multiple of that
+     * granularity (or fails if smaller than the granularity). Must be positive; defaults to 8 MB.
      */
     public abstract Builder setChunkSize(int chunkSize);
 

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

@@ -55,6 +55,12 @@ public abstract class ResumableUploadProgress {
    */
   public abstract @Nullable String getUploadUrl();
 
+  /**
+   * Returns the effective chunk size in bytes aligned to the server's chunk granularity, or {@code
+   * null} if session initiation is pending.
+   */
+  public abstract @Nullable Integer getEffectiveChunkSize();
+
   /** Returns the number of bytes confirmed as uploaded to the server so far. */
   public abstract long getBytesUploaded();
 
@@ -87,6 +93,8 @@ static Builder newBuilder() {
   abstract static class Builder {
     abstract Builder setUploadUrl(@Nullable String uploadUrl);
 
+    abstract Builder setEffectiveChunkSize(@Nullable Integer effectiveChunkSize);
+
     abstract Builder setBytesUploaded(long bytesUploaded);
 
     abstract Builder setState(String state);

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

@@ -116,10 +116,14 @@ ResumableUploadProgress getProgress() {
     }
   }
 
-  void onStarted(String uploadUrl) {
+  void onStarted(String uploadUrl, int effectiveChunkSize) {
     checkNotNull(uploadUrl, "uploadUrl must not be null");
     transition(
-        builder -> builder.setState(ResumableUploadProgress.STATE_STARTED).setUploadUrl(uploadUrl));
+        builder ->
+            builder
+                .setState(ResumableUploadProgress.STATE_STARTED)
+                .setUploadUrl(uploadUrl)
+                .setEffectiveChunkSize(effectiveChunkSize));
   }
 
   void onChunkUploaded(long bytesUploaded) {

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

@@ -1141,20 +1141,22 @@ void testProgressListener_prescribedStateTransitions() throws Exception {
     ResumableUploadFuture<String> future =
         callable.futureCall("resource-path", payloadOf("01234567890123456789"), options);
 
+    assertThat(future.getEffectiveChunkSize()).isNull();
+
     String url = "https://upload.url/progress-transitions";
     startFuture.set(ResumableUploadSession.newBuilder().setUploadUrl(url).build());
 
     assertThat(future.get()).isEqualTo("done");
 
     assertThat(listener.awaitTerminal())
         .containsExactly(
-            createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null),
-            createProgress(ResumableUploadProgress.STATE_STARTED, 0L, url),
-            createProgress(ResumableUploadProgress.STATE_UPLOADING, 8L, url),
-            createProgress(ResumableUploadProgress.STATE_RECOVERING, 8L, url),
-            createProgress(ResumableUploadProgress.STATE_OFFSET_RECEIVED, 8L, url),
-            createProgress(ResumableUploadProgress.STATE_UPLOADING, 16L, url),
-            createProgress(ResumableUploadProgress.STATE_FINALIZED, 20L, url))
+            createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null, null),
+            createProgress(ResumableUploadProgress.STATE_STARTED, 0L, url, 8),
+            createProgress(ResumableUploadProgress.STATE_UPLOADING, 8L, url, 8),
+            createProgress(ResumableUploadProgress.STATE_RECOVERING, 8L, url, 8),
+            createProgress(ResumableUploadProgress.STATE_OFFSET_RECEIVED, 8L, url, 8),
+            createProgress(ResumableUploadProgress.STATE_UPLOADING, 16L, url, 8),
+            createProgress(ResumableUploadProgress.STATE_FINALIZED, 20L, url, 8))
         .inOrder();
   }
 
@@ -1175,8 +1177,8 @@ void testProgressListener_uploadFailure_transitionsToFailed() throws Exception {
 
     assertThat(listener.awaitTerminal())
         .containsExactly(
-            createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null),
-            createProgress(ResumableUploadProgress.STATE_FAILED, 0L, null))
+            createProgress(ResumableUploadProgress.STATE_STARTING, 0L, null, null),
+            createProgress(ResumableUploadProgress.STATE_FAILED, 0L, null, null))
         .inOrder();
   }
 
@@ -1199,7 +1201,7 @@ void testProgressListener_cancel_transitionsToFailed() throws Exception {
     assertThat(receivedStatuses.get(receivedStatuses.size() - 1))
         .isEqualTo(
             createProgress(
-                ResumableUploadProgress.STATE_FAILED, 0L, "https://upload.url/cancel-progress"));
+                ResumableUploadProgress.STATE_FAILED, 0L, "https://upload.url/cancel-progress", 8));
   }
 
   @Test
@@ -1228,15 +1230,20 @@ void testProgressListener_queryFinal_transitionsToFinalized() throws Exception {
             createProgress(
                 ResumableUploadProgress.STATE_FINALIZED,
                 5L,
-                "https://upload.url/query-final-progress"));
+                "https://upload.url/query-final-progress",
+                8));
   }
 
   private static ResumableUploadProgress createProgress(
-      String state, long bytesUploaded, @Nullable String uploadUrl) {
+      String state,
+      long bytesUploaded,
+      @Nullable String uploadUrl,
+      @Nullable Integer effectiveChunkSize) {
     return ResumableUploadProgress.newBuilder()
         .setState(state)
         .setBytesUploaded(bytesUploaded)
         .setUploadUrl(uploadUrl)
+        .setEffectiveChunkSize(effectiveChunkSize)
         .build();
   }
 
@@ -1273,6 +1280,7 @@ void testActionableErrors_startFailure_preservesOriginalException() {
     ExecutionException ex = assertThrows(ExecutionException.class, future::get);
     assertThat(ex.getCause()).isSameInstanceAs(startError);
     assertThat(future.getUploadSessionUrl()).isNull();
+    assertThat(future.getEffectiveChunkSize()).isNull();
   }
 
   @Test
@@ -1298,9 +1306,53 @@ void testActionableErrors_preservesErrorDetailsAndCauseChain() {
     assertThat(ex.getCause()).isInstanceOf(ApiException.class);
     ApiException cause = (ApiException) ex.getCause();
     assertThat(cause.getMessage()).contains(sessionUrl);
+    assertThat(cause.getMessage()).contains("effective chunk size: 8");
     assertThat(cause.getCause()).isSameInstanceAs(transportCause);
     assertThat(cause.getErrorDetails()).isSameInstanceAs(errorDetails);
     assertThat(future.getUploadSessionUrl()).isEqualTo(sessionUrl);
+    assertThat(future.getEffectiveChunkSize()).isEqualTo(8);
+  }
+
+  @Test
+  void testAlignChunkSize() {
+    assertThat(ResumableUploadFutureImpl.alignChunkSize(8, 1L)).isEqualTo(8);
+    assertThat(ResumableUploadFutureImpl.alignChunkSize(8, 4L)).isEqualTo(8);
+    assertThat(ResumableUploadFutureImpl.alignChunkSize(10, 4L)).isEqualTo(8);
+    assertThrows(
+        IllegalArgumentException.class, () -> ResumableUploadFutureImpl.alignChunkSize(4, 6L));
+  }
+
+  @Test
+  void testChunkGranularity_alignsChunkSizeAndSurfacesEffectiveSize() throws Exception {
+    when(mockStartCallable.futureCall(any(), any()))
+        .thenReturn(
+            ApiFutures.immediateFuture(
+                ResumableUploadSession.newBuilder()
+                    .setUploadUrl("https://upload.url/granularity-align")
+                    .setChunkGranularity(4L)
+                    .build()));
+    when(mockChunkCallable.futureCall(any(ChunkUploadRequest.class), any()))
+        .thenReturn(
+            ApiFutures.immediateFuture(
+                ChunkUploadResponse.create(ResumableUploadStatus.ACTIVE, null)))
+        .thenReturn(
+            ApiFutures.immediateFuture(
+                ChunkUploadResponse.create(ResumableUploadStatus.FINAL, "done")));
+
+    // Configured chunkSize = 10, server granularity = 4 -> aligned down to 8
+    ResumableUploadOptions options = defaultSettings.toBuilder().setChunkSize(10).build();
+
+    ResumableUploadFuture<String> future =
+        callable.futureCall("resource-path", payloadOf("0123456789ab"), options);
+
+    assertThat(future.get()).isEqualTo("done");
+    assertThat(future.getEffectiveChunkSize()).isEqualTo(8);
+
+    ArgumentCaptor<ChunkUploadRequest> captor = ArgumentCaptor.forClass(ChunkUploadRequest.class);
+    verify(mockChunkCallable, times(2)).futureCall(captor.capture(), any());
+    List<ChunkUploadRequest> requests = captor.getAllValues();
+    assertChunk(requests.get(0), 0, 8, false);
+    assertChunk(requests.get(1), 8, 4, true);
   }
 
   @Test