Add more tests and address comments

Author: zoewangg

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

@@ -20,10 +20,38 @@
 import org.reactivestreams.Subscriber;
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.core.internal.async.SplittingPublisher;
+import software.amazon.awssdk.utils.Validate;
 
 /**
- * An {@link AsyncRequestBody} decorator that can be split into buffered sub {@link AsyncRequestBody}s. Each sub
- * {@link AsyncRequestBody} can be retried/resubscribed if all data has been successfully been published to first subscriber.
+ * An {@link AsyncRequestBody} decorator that enables splitting into retryable sub-request bodies.
+ *
+ * <p>This wrapper allows any {@link AsyncRequestBody} to be split into multiple parts where each part
+ * can be retried independently. When split, each sub-body buffers its portion of data, enabling
+ * resubscription if a retry is needed (e.g., due to network failures or service errors).</p>
+ *
+ * <p><b>Usage Example:</b></p>
+ * <pre>{@code
+ * AsyncRequestBody originalBody = AsyncRequestBody.fromString("Hello World");
+ * BufferedSplittableAsyncRequestBody retryableBody =
+ *     BufferedSplittableAsyncRequestBody.create(originalBody);
+ *
+ * AsyncRequestBodySplitConfiguration config = AsyncRequestBodySplitConfiguration.builder()
+ *     .chunkSizeInBytes(1024)
+ *     .bufferSizeInBytes(2048)
+ *     .build();
+ *
+ * SdkPublisher<ClosableAsyncRequestBody> parts = retryableBody.splitClosable(config);
+ * }</pre>
+ *
+ * <p><b>Performance Considerations:</b></p>
+ * <p>This implementation buffers data in memory to enable retries, but memory usage is controlled by
+ * the {@code bufferSizeInBytes} configuration. However, this buffering limits the ability to request
+ * more data from the original AsyncRequestBody until buffered data is consumed (i.e., when subscribers
+ * closes sub-body), which may increase latency compared to non-buffered implementations.
+ *
+ * @see AsyncRequestBody
+ * @see AsyncRequestBodySplitConfiguration
+ * @see ClosableAsyncRequestBody
  */
 @SdkPublicApi
 public final class BufferedSplittableAsyncRequestBody implements AsyncRequestBody {
@@ -33,7 +61,15 @@ private BufferedSplittableAsyncRequestBody(AsyncRequestBody delegate) {
         this.delegate = delegate;
     }
 
+    /**
+     * Creates a new {@link BufferedSplittableAsyncRequestBody} that wraps the provided {@link AsyncRequestBody}.
+     *
+     * @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);
     }
 
@@ -42,6 +78,22 @@ public Optional<Long> contentLength() {
         return delegate.contentLength();
     }
 
+    /**
+     * Splits this request body into multiple retryable parts based on the provided configuration.
+     *
+     * <p>Each part returned by the publisher will be a {@link ClosableAsyncRequestBody} that buffers
+     * its portion of data, enabling resubscription for retry scenarios. This is the key difference from non-buffered splitting -
+     * each part can be safely retried without data loss.
+     *
+     * <p>The splitting process respects the chunk size and buffer size specified in the configuration
+     * to optimize memory usage.
+     *
+     * <p>The subscriber MUST close each {@link ClosableAsyncRequestBody} to ensure resource is released
+     *
+     * @param splitConfiguration configuration specifying how to split the request body
+     * @return a publisher that emits retryable closable request body parts
+     * @see AsyncRequestBodySplitConfiguration
+     */
     @Override
     public SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
         return new SplittingPublisher(this, splitConfiguration, true);

core/sdk-core/src/main/java/software/amazon/awssdk/core/async/ClosableAsyncRequestBody.java renamed to core/sdk-core/src/main/java/software/amazon/awssdk/core/async/CloseableAsyncRequestBody.java

@@ -22,5 +22,5 @@
  * An extension of {@link AsyncRequestBody} that is closable.
  */
 @SdkPublicApi
-public interface ClosableAsyncRequestBody extends AsyncRequestBody, SdkAutoCloseable {
+public interface CloseableAsyncRequestBody extends AsyncRequestBody, SdkAutoCloseable {
 }

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

@@ -23,7 +23,7 @@
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncRequestBodySplitConfiguration;
-import software.amazon.awssdk.core.async.ClosableAsyncRequestBody;
+import software.amazon.awssdk.core.async.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
@@ -78,14 +78,14 @@ public SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfig
         }
 
         @Override
-        public SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
-            return delegate.splitClosable(splitConfiguration);
+        public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+            return delegate.splitCloseable(splitConfiguration);
         }
 
         @Override
-        public SdkPublisher<ClosableAsyncRequestBody> splitClosable(
+        public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(
             Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
-            return delegate.splitClosable(splitConfiguration);
+            return delegate.splitCloseable(splitConfiguration);
         }
 
         @Override

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

@@ -34,7 +34,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncRequestBodySplitConfiguration;
-import software.amazon.awssdk.core.async.ClosableAsyncRequestBody;
+import software.amazon.awssdk.core.async.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.internal.util.Mimetype;
 import software.amazon.awssdk.core.internal.util.NoopSubscription;
@@ -88,7 +88,7 @@ public SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration s
     }
 
     @Override
-    public SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+    public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
         return split(splitConfiguration).map(body -> new ClosableAsyncRequestBodyWrapper(body));
     }
 
@@ -443,7 +443,7 @@ private static AsynchronousFileChannel openInputChannel(Path path) throws IOExce
         return AsynchronousFileChannel.open(path, StandardOpenOption.READ);
     }
 
-    private static class ClosableAsyncRequestBodyWrapper implements ClosableAsyncRequestBody {
+    private static class ClosableAsyncRequestBodyWrapper implements CloseableAsyncRequestBody {
         private final AsyncRequestBody delegate;
 
         ClosableAsyncRequestBodyWrapper(AsyncRequestBody body) {

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

@@ -62,11 +62,11 @@ public Optional<Long> contentLength() {
 
     public void send(ByteBuffer data) {
         log.debug(() -> String.format("Sending bytebuffer %s to part %d", data, partNumber));
-        int length = data.remaining();
+        long length = data.remaining();
         bufferedLength += length;
-        onNumBytesReceived.accept((long) length);
+        onNumBytesReceived.accept(length);
         delegate.send(data).whenComplete((r, t) -> {
-            onNumBytesConsumed.accept((long) length);
+            onNumBytesConsumed.accept(length);
             if (t != null) {
                 error(t);
             }
@@ -92,11 +92,6 @@ public long receivedBytesLength() {
         return bufferedLength;
     }
 
-    @Override
-    public boolean contentLengthKnown() {
-        return contentLengthKnown;
-    }
-
     @Override
     public int partNumber() {
         return partNumber;
@@ -113,8 +108,10 @@ public void subscribe(Subscriber<? super ByteBuffer> s) {
         } else {
             s.onSubscribe(new NoopSubscription(s));
             s.onError(NonRetryableException.create(
-                "A retry was attempted, but the provided source AsyncRequestBody does not "
-                + "support splitting to retryable AsyncRequestBody. Consider using BufferedSplittableAsyncRequestBody."));
+                "Multiple subscriptions detected. This could happen due to a retry attempt. The AsyncRequestBody implementation"
+                + " provided does not support splitting to retryable/resubscribable AsyncRequestBody. If you need retry "
+                + "capability or multiple subscriptions, consider using BufferedSplittableAsyncRequestBody to wrap your "
+                + "AsyncRequestBody."));
         }
     }
 

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

@@ -75,10 +75,10 @@ public Optional<Long> contentLength() {
     @Override
     public void send(ByteBuffer data) {
         log.trace(() -> String.format("Sending bytebuffer %s to part number %d", data, partNumber));
-        int length = data.remaining();
+        long length = data.remaining();
         bufferedLength += length;
 
-        onNumBytesReceived.accept((long) length);
+        onNumBytesReceived.accept(length);
         delegate.send(data.asReadOnlyBuffer()).whenComplete((r, t) -> {
             if (t != null) {
                 delegate.error(t);
@@ -94,7 +94,9 @@ public void complete() {
         log.debug(() -> "Received complete() for part number: " + partNumber);
         // ByteBuffersAsyncRequestBody MUST be created before we complete the current
         // request because retry may happen right after
-        bufferedAsyncRequestBody = ByteBuffersAsyncRequestBody.of(buffers, bufferedLength);
+        synchronized (buffersLock) {
+            bufferedAsyncRequestBody = ByteBuffersAsyncRequestBody.of(buffers, bufferedLength);
+        }
         delegate.complete().exceptionally(e -> {
             delegate.error(e);
             return null;
@@ -121,7 +123,7 @@ public void subscribe(Subscriber<? super ByteBuffer> s) {
             if (bufferedAsyncRequestBody == null) {
                 s.onSubscribe(new NoopSubscription(s));
                 s.onError(NonRetryableException.create(
-                    "A retry was attempted, but data is not buffered successfully for retry, partNumber " + partNumber));
+                    "A retry was attempted, but data is not buffered successfully for retry for partNumber: " + partNumber));
                 return;
             }
             bufferedAsyncRequestBody.subscribe(s);
@@ -139,19 +141,14 @@ public void close() {
                     buffers = null;
                 }
                 bufferedAsyncRequestBody.close();
-                log.debug(() -> "requesting data after closing" + partNumber);
+                bufferedAsyncRequestBody = null;
             }
         } catch (Throwable e) {
             log.warn(() -> String.format("Unexpected error thrown from cleaning up AsyncRequestBody for part number %d, "
                                          + "resource may be leaked", partNumber));
         }
     }
 
-    @Override
-    public boolean contentLengthKnown() {
-        return contentLengthKnown;
-    }
-
     @Override
     public int partNumber() {
         return partNumber;

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

@@ -24,7 +24,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncRequestBodySplitConfiguration;
-import software.amazon.awssdk.core.async.ClosableAsyncRequestBody;
+import software.amazon.awssdk.core.async.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
@@ -38,43 +38,49 @@
  * Otherwise, it is sent after the entire content for that chunk is buffered. This is required to get content length.
  */
 @SdkInternalApi
-public class SplittingPublisher implements SdkPublisher<ClosableAsyncRequestBody> {
+public class SplittingPublisher implements SdkPublisher<CloseableAsyncRequestBody> {
     private static final Logger log = Logger.loggerFor(SplittingPublisher.class);
     private final AsyncRequestBody upstreamPublisher;
     private final SplittingSubscriber splittingSubscriber;
-    private final SimplePublisher<ClosableAsyncRequestBody> downstreamPublisher = new SimplePublisher<>();
+    private final SimplePublisher<CloseableAsyncRequestBody> downstreamPublisher = new SimplePublisher<>();
     private final long chunkSizeInBytes;
     private final long bufferSizeInBytes;
-    private final boolean enableRetryableSubAsyncRequestBody;
+    private final boolean retryableSubAsyncRequestBodyEnabled;
     private final AtomicBoolean currentBodySent = new AtomicBoolean(false);
     private final String sourceBodyName;
 
-    public SplittingPublisher(AsyncRequestBody asyncRequestBody,
-                              AsyncRequestBodySplitConfiguration splitConfiguration,
-                              boolean enableRetryableSubAsyncRequestBody) {
-        this.upstreamPublisher = Validate.paramNotNull(asyncRequestBody, "asyncRequestBody");
-        Validate.notNull(splitConfiguration, "splitConfiguration");
-        this.chunkSizeInBytes = splitConfiguration.chunkSizeInBytes() == null ?
+    private SplittingPublisher(Builder builder) {
+        this.upstreamPublisher = Validate.paramNotNull(builder.asyncRequestBody, "asyncRequestBody");
+        Validate.notNull(builder.splitConfiguration, "splitConfiguration");
+        this.chunkSizeInBytes = builder.splitConfiguration.chunkSizeInBytes() == null ?
                                 AsyncRequestBodySplitConfiguration.defaultConfiguration().chunkSizeInBytes() :
-                                splitConfiguration.chunkSizeInBytes();
+                                builder.splitConfiguration.chunkSizeInBytes();
 
-        this.bufferSizeInBytes = splitConfiguration.bufferSizeInBytes() == null ?
+        this.bufferSizeInBytes = builder.splitConfiguration.bufferSizeInBytes() == null ?
                                  AsyncRequestBodySplitConfiguration.defaultConfiguration().bufferSizeInBytes() :
-                                 splitConfiguration.bufferSizeInBytes();
+                                 builder.splitConfiguration.bufferSizeInBytes();
 
         this.splittingSubscriber = new SplittingSubscriber(upstreamPublisher.contentLength().orElse(null));
 
-        this.enableRetryableSubAsyncRequestBody = enableRetryableSubAsyncRequestBody;
-        this.sourceBodyName = asyncRequestBody.body();
+        this.retryableSubAsyncRequestBodyEnabled = Validate.paramNotNull(builder.retryableSubAsyncRequestBodyEnabled,
+                                                                         "retryableSubAsyncRequestBodyEnabled");
+        this.sourceBodyName = builder.asyncRequestBody.body();
         if (!upstreamPublisher.contentLength().isPresent()) {
             Validate.isTrue(bufferSizeInBytes >= chunkSizeInBytes,
                             "bufferSizeInBytes must be larger than or equal to " +
                             "chunkSizeInBytes if the content length is unknown");
         }
     }
 
+    /**
+     * Returns a newly initialized builder object for a {@link SplittingPublisher}
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
     @Override
-    public void subscribe(Subscriber<? super ClosableAsyncRequestBody> downstreamSubscriber) {
+    public void subscribe(Subscriber<? super CloseableAsyncRequestBody> downstreamSubscriber) {
         downstreamPublisher.subscribe(downstreamSubscriber);
         upstreamPublisher.subscribe(splittingSubscriber);
     }
@@ -123,7 +129,7 @@ private SubAsyncRequestBody initializeNextDownstreamBody(boolean contentLengthKn
                     .sourceBodyName(sourceBodyName)
                     .build();
 
-            if (enableRetryableSubAsyncRequestBody) {
+            if (retryableSubAsyncRequestBodyEnabled) {
                 body = new RetryableSubAsyncRequestBody(config);
             } else {
                 body = new NonRetryableSubAsyncRequestBody(config);
@@ -206,15 +212,16 @@ private int amountRemainingInChunk() {
 
         private void completeCurrentBody() {
             log.debug(() -> "completeCurrentBody for part " + currentBody.partNumber());
-            // For unknown content length, we always create a new DownstreamBody because we don't know if there is data
-            // left or not, so we need to only send the body if there is actually data
             long bufferedLength = currentBody.receivedBytesLength();
+            // For unknown content length, we always create a new DownstreamBody once the current one is sent
+            // because we don't know if there is data
+            // left or not, so we need to check the length and only send the body if there is actually data
             if (bufferedLength == 0) {
                 return;
             }
 
             Long totalLength = currentBody.maxLength();
-            if (currentBody.contentLengthKnown() && totalLength != bufferedLength) {
+            if (upstreamSize != null && totalLength != bufferedLength) {
                 upstreamSubscription.cancel();
                 downstreamPublisher.error(new IllegalStateException(
                     String.format("Content length of buffered data mismatches "
@@ -294,4 +301,44 @@ private void addDataBuffered(long length) {
             }
         }
     }
+
+    public static final class Builder {
+        private AsyncRequestBody asyncRequestBody;
+        private AsyncRequestBodySplitConfiguration splitConfiguration;
+        private Boolean retryableSubAsyncRequestBodyEnabled;
+
+        private Builder() {
+        }
+
+        /**
+         * Sets the AsyncRequestBody to be split.
+         */
+        public Builder asyncRequestBody(AsyncRequestBody asyncRequestBody) {
+            this.asyncRequestBody = asyncRequestBody;
+            return this;
+        }
+
+        /**
+         * Sets the split configuration.
+         */
+        public Builder splitConfiguration(AsyncRequestBodySplitConfiguration splitConfiguration) {
+            this.splitConfiguration = splitConfiguration;
+            return this;
+        }
+
+        /**
+         * Sets whether to enable retryable sub async request bodies.
+         */
+        public Builder retryableSubAsyncRequestBodyEnabled(Boolean retryableSubAsyncRequestBodyEnabled) {
+            this.retryableSubAsyncRequestBodyEnabled = retryableSubAsyncRequestBodyEnabled;
+            return this;
+        }
+
+        /**
+         * Builds a {@link SplittingPublisher} object based on the values held by this builder.
+         */
+        public SplittingPublisher build() {
+            return new SplittingPublisher(this);
+        }
+    }
 }