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

Author: whowes

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

@@ -108,7 +108,7 @@ static <RequestT> UnaryCallable<RequestT, ResumableUploadSession> create(
   /** A listener that parses HTTP response headers to produce a {@link ResumableUploadSession}. */
   private static class StartUploadResponseListener extends HttpJsonClientCall.Listener<String> {
     private final ResumableUploadHttpJsonFuture<ResumableUploadSession> future;
-    private long chunkGranularity = 1L;
+    private int chunkGranularity = 1;
     @Nullable private String uploadUrl;
     @Nullable private Throwable headerParsingException;
 
@@ -132,7 +132,7 @@ public void onHeaders(HttpJsonMetadata responseHeaders) {
       }
 
       try {
-        long parsed = Long.parseLong(granularityStr);
+        int parsed = Integer.parseInt(granularityStr);
         if (parsed <= 0) {
           this.headerParsingException =
               ApiExceptionFactory.createException(

sdk-platform-java/gax-java/gax-httpjson/src/test/java/com/google/api/gax/httpjson/HttpJsonResumableUploadClientTest.java

@@ -105,7 +105,7 @@ void startUpload_validHeaders_returnsSession() {
     ResumableUploadSession session = client.startUploadCallable().call(request);
 
     assertThat(session.getUploadUrl()).isEqualTo(TEST_UPLOAD_URL);
-    assertThat(session.getChunkGranularity()).isEqualTo(262144L);
+    assertThat(session.getChunkGranularity()).isEqualTo(262144);
   }
 
   @Test
@@ -123,7 +123,7 @@ void startUpload_caseInsensitiveHeaders_returnsSession() {
 
     assertThat(session.getUploadUrl())
         .isEqualTo("https://test.googleapis.com/upload/session/case-insensitive");
-    assertThat(session.getChunkGranularity()).isEqualTo(524288L);
+    assertThat(session.getChunkGranularity()).isEqualTo(524288);
   }
 
   @Test

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

@@ -41,7 +41,7 @@
 @AutoValue
 public abstract class ResumableUploadSession {
 
-  private static final long DEFAULT_CHUNK_GRANULARITY = 1L;
+  private static final int DEFAULT_CHUNK_GRANULARITY = 1;
 
   /** Returns the server-provided URL to which data uploads are directed. */
   public abstract String getUploadUrl();
@@ -56,7 +56,7 @@ public abstract class ResumableUploadSession {
    *
    * @return the chunk granularity in bytes
    */
-  public abstract long getChunkGranularity();
+  public abstract int getChunkGranularity();
 
   public abstract Builder toBuilder();
 
@@ -69,7 +69,7 @@ public static Builder newBuilder() {
   public abstract static class Builder {
     public abstract Builder setUploadUrl(String uploadUrl);
 
-    public abstract Builder setChunkGranularity(long chunkGranularity);
+    public abstract Builder setChunkGranularity(int chunkGranularity);
 
     public abstract ResumableUploadSession build();
   }

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,19 @@
  * @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. */
+  /**
+   * Returns the upload session URL, or {@code null} if session initiation is in progress or did not
+   * complete successfully.
+   */
   @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 or did not complete successfully.
+   */
+  @Nullable Integer getEffectiveChunkSize();
 }

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

@@ -85,7 +85,12 @@ final class ResumableUploadFutureImpl<ResponseT> implements ResumableUploadFutur
   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 @@ private void start(InputStreamSupplier payloadSupplier) throws IOException {
           @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 @@ public void onSuccess(ResumableUploadSession session) {
               if (inFlightFuture == null) {
                 return;
               }
+              effectiveChunkSize = chunkSize;
               uploadSessionUrl = sessionUrl;
               inFlightFuture = uploadFuture;
             }
@@ -213,7 +226,7 @@ public void onFailure(Throwable t) {
 
   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 void fail(Throwable t) {
   }
 
   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,16 @@ private ApiException toTerminalException(Throwable t) {
         message, t, statusCodeOf(StatusCode.Code.UNKNOWN), false);
   }
 
+  static int alignChunkSize(int configuredChunkSize, int chunkGranularity) {
+    checkArgument(configuredChunkSize > 0, "configuredChunkSize must be > 0");
+    checkArgument(chunkGranularity > 0, "chunkGranularity must be > 0");
+    if (configuredChunkSize <= chunkGranularity) {
+      return chunkGranularity;
+    }
+    int unalignedRemainder = configuredChunkSize % chunkGranularity;
+    return configuredChunkSize - unalignedRemainder;
+  }
+
   private void closePayload(@Nullable Throwable primaryException) {
     if (payload == null) {
       return;
@@ -287,7 +318,16 @@ private void closePayload(@Nullable Throwable primaryException) {
 
   @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 raised to the granularity if smaller). 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/resumable/ResumableUploadSessionTest.java

@@ -43,12 +43,12 @@ void build_withUploadUrl_setsDefaultChunkGranularity() {
         ResumableUploadSession.newBuilder().setUploadUrl(UPLOAD_URL).build();
 
     assertThat(session.getUploadUrl()).isEqualTo(UPLOAD_URL);
-    assertThat(session.getChunkGranularity()).isEqualTo(1L);
+    assertThat(session.getChunkGranularity()).isEqualTo(1);
   }
 
   @Test
   void build_withExplicitChunkGranularity_preservesGranularity() {
-    long customChunkGranularity = 256 * 1024L;
+    int customChunkGranularity = 256 * 1024;
 
     ResumableUploadSession session =
         ResumableUploadSession.newBuilder()