Add more tests and address comments

Author: zoewangg

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

@@ -516,11 +516,16 @@ static AsyncRequestBody empty() {
      */
     default SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
-        return new SplittingPublisher(this, splitConfiguration, false).map(r -> r);
+        return SplittingPublisher.builder()
+                .asyncRequestBody(this)
+                .splitConfiguration(splitConfiguration)
+                .retryableSubAsyncRequestBodyEnabled(false)
+                .build()
+                .map(r -> r);
     }
 
     /**
-     * Converts this {@link AsyncRequestBody} to a publisher of {@link ClosableAsyncRequestBody}s, each of which publishes
+     * Converts this {@link AsyncRequestBody} to a publisher of {@link CloseableAsyncRequestBody}s, each of which publishes
      * specific portion of the original data, based on the provided {@link AsyncRequestBodySplitConfiguration}. The default chunk
      * size is 2MB and the default buffer size is 8MB.
      *
@@ -529,14 +534,18 @@ default SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration
      * vary in different implementations.
      *
      * <p>
-     * Caller is responsible for closing {@link ClosableAsyncRequestBody} when it is ready to be disposed to release any
+     * Caller is responsible for closing {@link CloseableAsyncRequestBody} when it is ready to be disposed to release any
      * resources.
      *
      * @see AsyncRequestBodySplitConfiguration
      */
-    default SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+    default SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
-        return new SplittingPublisher(this, splitConfiguration, false);
+        return SplittingPublisher.builder()
+                .asyncRequestBody(this)
+                .splitConfiguration(splitConfiguration)
+                .retryableSubAsyncRequestBodyEnabled(false)
+                .build();
     }
 
     /**
@@ -554,12 +563,12 @@ default SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfi
      * This is a convenience method that passes an instance of the {@link AsyncRequestBodySplitConfiguration} builder,
      * avoiding the need to create one manually via {@link AsyncRequestBodySplitConfiguration#builder()}.
      *
-     * @see #splitClosable(Consumer)
+     * @see #splitCloseable(Consumer)
      */
-    default SdkPublisher<ClosableAsyncRequestBody> splitClosable(
+    default SdkPublisher<CloseableAsyncRequestBody> splitCloseable(
         Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
-        return splitClosable(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
+        return splitCloseable(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
     }
 
     @SdkProtectedApi

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

@@ -20,10 +20,43 @@
 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>Retry Requirements:</b></p>
+ * <p>Retry is only possible if all the data has been successfully buffered during the first subscription.
+ * If the first subscriber fails to consume all the data (e.g., due to early cancellation or errors),
+ * subsequent retry attempts will fail since the complete data set is not available for resubscription.</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 CloseableAsyncRequestBody
  */
 @SdkPublicApi
 public final class BufferedSplittableAsyncRequestBody implements AsyncRequestBody {
@@ -33,7 +66,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,9 +83,29 @@ 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 CloseableAsyncRequestBody} 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 CloseableAsyncRequestBody} 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);
+    public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+        return SplittingPublisher.builder()
+                .asyncRequestBody(this)
+                .splitConfiguration(splitConfiguration)
+                .retryableSubAsyncRequestBodyEnabled(true)
+                .build();
     }
 
     @Override

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 subscribers 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;