Add fullBufferingEnabled opt-in option to BufferedSplittableAsyncRequestBody

Author: alextwoods

core/sdk-core/src/main/java/software/amazon/awssdk/core/async/BufferedSplittableAsyncRequestBody.java

@@ -54,21 +54,46 @@
 @SdkPublicApi
 public final class BufferedSplittableAsyncRequestBody implements AsyncRequestBody {
     private final AsyncRequestBody delegate;
+    private final boolean fullBufferingEnabled;
 
-    private BufferedSplittableAsyncRequestBody(AsyncRequestBody delegate) {
+    private BufferedSplittableAsyncRequestBody(AsyncRequestBody delegate, boolean fullBufferingEnabled) {
         this.delegate = delegate;
+        this.fullBufferingEnabled = fullBufferingEnabled;
     }
 
     /**
      * Creates a new {@link BufferedSplittableAsyncRequestBody} that wraps the provided {@link AsyncRequestBody}.
      *
+     * <p>Full buffering is disabled by default. Each part is sent to the downstream subscriber immediately
+     * upon initialization in the known-content-length path (existing behavior).
+     *
      * @param delegate the {@link AsyncRequestBody} to wrap and make retryable. Must not be null.
      * @return a new {@link BufferedSplittableAsyncRequestBody} instance
      * @throws NullPointerException if delegate is null
      */
     public static BufferedSplittableAsyncRequestBody create(AsyncRequestBody delegate) {
         Validate.paramNotNull(delegate, "delegate");
-        return new BufferedSplittableAsyncRequestBody(delegate);
+        return new BufferedSplittableAsyncRequestBody(delegate, false);
+    }
+
+    /**
+     * Creates a new {@link BufferedSplittableAsyncRequestBody} that wraps the provided {@link AsyncRequestBody},
+     * with an option to enable full buffering before sending parts downstream.
+     *
+     * <p>When {@code fullBufferingEnabled} is {@code true}, each part is fully buffered before being sent to the
+     * downstream subscriber. This guarantees that the retry buffer is always populated before the HTTP layer
+     * subscribes, making per-part retry deterministically successful for slow streaming sources (e.g., SFTP).
+     *
+     * <p>When {@code fullBufferingEnabled} is {@code false}, behavior is identical to {@link #create(AsyncRequestBody)}.
+     *
+     * @param delegate the {@link AsyncRequestBody} to wrap and make retryable. Must not be null.
+     * @param fullBufferingEnabled whether to enable full buffering before sending parts downstream
+     * @return a new {@link BufferedSplittableAsyncRequestBody} instance
+     * @throws NullPointerException if delegate is null
+     */
+    public static BufferedSplittableAsyncRequestBody create(AsyncRequestBody delegate, boolean fullBufferingEnabled) {
+        Validate.paramNotNull(delegate, "delegate");
+        return new BufferedSplittableAsyncRequestBody(delegate, fullBufferingEnabled);
     }
 
     @Override
@@ -98,6 +123,7 @@ public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySp
                 .asyncRequestBody(this)
                 .splitConfiguration(splitConfiguration)
                 .retryableSubAsyncRequestBodyEnabled(true)
+                .fullBufferingEnabled(fullBufferingEnabled)
                 .build();
     }
 

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/SplittingPublisher.java

@@ -46,6 +46,7 @@ public class SplittingPublisher implements SdkPublisher<CloseableAsyncRequestBod
     private final long chunkSizeInBytes;
     private final long bufferSizeInBytes;
     private final boolean retryableSubAsyncRequestBodyEnabled;
+    private final boolean fullBufferingEnabled;
     private final AtomicBoolean currentBodySent = new AtomicBoolean(false);
     private final String sourceBodyName;
 
@@ -64,6 +65,7 @@ private SplittingPublisher(Builder builder) {
 
         this.retryableSubAsyncRequestBodyEnabled = Validate.paramNotNull(builder.retryableSubAsyncRequestBodyEnabled,
                                                                          "retryableSubAsyncRequestBodyEnabled");
+        this.fullBufferingEnabled = builder.fullBufferingEnabled;
         this.sourceBodyName = builder.asyncRequestBody.body();
         if (!upstreamPublisher.contentLength().isPresent()) {
             Validate.isTrue(bufferSizeInBytes >= chunkSizeInBytes,
@@ -136,7 +138,7 @@ private SubAsyncRequestBody initializeNextDownstreamBody(boolean contentLengthKn
             }
 
             currentBodySent.set(false);
-            if (contentLengthKnown) {
+            if (contentLengthKnown && !fullBufferingEnabled) {
                 sendCurrentBody(body);
             }
             return body;
@@ -234,7 +236,7 @@ private void completeCurrentBody() {
 
             // Current body could be completed in either onNext or onComplete, so we need to guard against sending the last body
             // twice.
-            if (upstreamSize == null && currentBodySent.compareAndSet(false, true)) {
+            if ((upstreamSize == null || fullBufferingEnabled) && currentBodySent.compareAndSet(false, true)) {
                 sendCurrentBody(currentBody);
             }
         }
@@ -307,6 +309,7 @@ public static final class Builder {
         private AsyncRequestBody asyncRequestBody;
         private AsyncRequestBodySplitConfiguration splitConfiguration;
         private Boolean retryableSubAsyncRequestBodyEnabled;
+        private boolean fullBufferingEnabled = false;
 
         private Builder() {
         }
@@ -335,6 +338,16 @@ public Builder retryableSubAsyncRequestBodyEnabled(Boolean retryableSubAsyncRequ
             return this;
         }
 
+        /**
+         * Sets whether to enable full buffering before sending parts downstream.
+         * When enabled, parts are only sent to the downstream subscriber after
+         * all data for that part has been received and complete() has been called.
+         */
+        public Builder fullBufferingEnabled(boolean fullBufferingEnabled) {
+            this.fullBufferingEnabled = fullBufferingEnabled;
+            return this;
+        }
+
         /**
          * Builds a {@link SplittingPublisher} object based on the values held by this builder.
          */