Stream retry support part 2: Introduce a new split method in AsyncRequestBody that returns an SdkP… (#6346)

* Introduce a new split method in AsyncRequestBody that returns an SdkPublisher of ClosableAsyncRequestBody and use it in s3 multipart client

* Fix build and add more tests

* Refactoring based on API surface are review

* Add more tests and address comments

* Deprecate split, update documentation and add changelog entries

Author: zoewangg

.changes/next-release/deprecation-AWSSDKforJavav2-6069419.json

@@ -0,0 +1,6 @@
+{
+    "type": "deprecation",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Deprecate `AsyncRequestBody#split` in favor of `AsyncRequestBody#splitCloseable` that takes the same input but returns `SdkPublisher<CloseableAsyncRequestBody>`"
+}

.changes/next-release/feature-AWSSDKforJavav2-500254f.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Add `AsyncRequestBody#splitCloseable` API that returns a Publisher of `ClosableAsyncRequestBody`"
+}

.changes/next-release/feature-AWSSDKforJavav2-72c5f55.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Introduce CloseableAsyncRequestBody interface that extends both AsyncRequestBody and SdkAutoClosable interfaces"
+}

.changes/next-release/feature-AWSSDKforJavav2-ed2c57a.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "AWS SDK for Java v2",
+    "contributor": "",
+    "description": "Introduce BufferedSplittableAsyncRequestBody that enables splitting into retryable sub-request bodies."
+}

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

@@ -500,7 +500,6 @@ static AsyncRequestBody empty() {
         return fromBytes(new byte[0]);
     }
 
-
     /**
      * Converts this {@link AsyncRequestBody} to a publisher of {@link AsyncRequestBody}s, each of which publishes a specific
      * portion of the original data, based on the provided {@link AsyncRequestBodySplitConfiguration}. The default chunk size
@@ -513,25 +512,75 @@ static AsyncRequestBody empty() {
      * than or equal to {@code chunkSizeInBytes}. Note that this behavior may be different if a specific implementation of this
      * interface overrides this method.
      *
-     * @see AsyncRequestBodySplitConfiguration
+     * @deprecated use {@link #splitCloseable(AsyncRequestBodySplitConfiguration)} instead.
      */
+    @Deprecated
     default SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
+        return SplittingPublisher.builder()
+                .asyncRequestBody(this)
+                .splitConfiguration(splitConfiguration)
+                .retryableSubAsyncRequestBodyEnabled(false)
+                .build()
+                .map(r -> r);
+    }
 
-        return new SplittingPublisher(this, splitConfiguration);
+    /**
+     * 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.
+     *
+     * <p>
+     * The default implementation behaves the same as {@link #split(AsyncRequestBodySplitConfiguration)}. This behavior may
+     * vary in different implementations.
+     *
+     * <p>
+     * Caller is responsible for closing {@link CloseableAsyncRequestBody} when it is ready to be disposed to release any
+     * resources.
+     *
+     * <p><b>Note:</b> This method is primarily intended for use by AWS SDK high-level libraries and internal components.
+     * SDK customers should typically use higher-level APIs provided by service clients rather than calling this method directly.
+     *
+     * @see #splitCloseable(Consumer)
+     * @see AsyncRequestBodySplitConfiguration
+     */
+    default SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+        Validate.notNull(splitConfiguration, "splitConfiguration");
+        return SplittingPublisher.builder()
+                .asyncRequestBody(this)
+                .splitConfiguration(splitConfiguration)
+                .retryableSubAsyncRequestBodyEnabled(false)
+                .build();
     }
 
     /**
      * 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 #split(AsyncRequestBodySplitConfiguration)
+     * @deprecated use {@link #splitCloseable(Consumer)} instead
      */
+    @Deprecated
     default SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
         return split(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
     }
 
+    /**
+     * 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()}.
+     *
+     * <p><b>Note:</b> This method is primarily intended for use by AWS SDK high-level libraries and internal components.
+     * SDK customers should typically use higher-level APIs provided by service clients rather than calling this method directly.
+     *
+     * @see #splitCloseable(AsyncRequestBodySplitConfiguration)
+     */
+    default SdkPublisher<CloseableAsyncRequestBody> splitCloseable(
+        Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
+        Validate.notNull(splitConfiguration, "splitConfiguration");
+        return splitCloseable(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
+    }
+
     @SdkProtectedApi
     enum BodyType {
         FILE("File", "f"),

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

@@ -0,0 +1,113 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.async;
+
+import java.nio.ByteBuffer;
+import java.util.Optional;
+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 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>
+ * {@snippet :
+ * AsyncRequestBody originalBody = AsyncRequestBody.fromString("Hello World");
+ * BufferedSplittableAsyncRequestBody retryableBody =
+ *     BufferedSplittableAsyncRequestBody.create(originalBody);
+ * }
+ *
+ * <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 {
+    private final AsyncRequestBody delegate;
+
+    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);
+    }
+
+    @Override
+    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<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+        return SplittingPublisher.builder()
+                .asyncRequestBody(this)
+                .splitConfiguration(splitConfiguration)
+                .retryableSubAsyncRequestBodyEnabled(true)
+                .build();
+    }
+
+    @Override
+    public void subscribe(Subscriber<? super ByteBuffer> s) {
+        delegate.subscribe(s);
+    }
+
+    @Override
+    public String body() {
+        return delegate.body();
+    }
+}

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

@@ -0,0 +1,26 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.core.async;
+
+import software.amazon.awssdk.annotations.SdkPublicApi;
+import software.amazon.awssdk.utils.SdkAutoCloseable;
+
+/**
+ * An extension of {@link AsyncRequestBody} that is closable.
+ */
+@SdkPublicApi
+public interface CloseableAsyncRequestBody extends AsyncRequestBody, SdkAutoCloseable {
+}

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

@@ -23,6 +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.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
@@ -76,6 +77,17 @@ public SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfig
             return delegate.split(splitConfiguration);
         }
 
+        @Override
+        public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+            return delegate.splitCloseable(splitConfiguration);
+        }
+
+        @Override
+        public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(
+            Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
+            return delegate.splitCloseable(splitConfiguration);
+        }
+
         @Override
         public void subscribe(Subscriber<? super ByteBuffer> s) {
             invoke(() -> listener.publisherSubscribe(s), "publisherSubscribe");

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

@@ -76,7 +76,9 @@ public final class ByteBuffersAsyncRequestBody implements AsyncRequestBody, SdkA
     private final Object lock = new Object();
     private boolean closed;
 
-    private ByteBuffersAsyncRequestBody(String mimetype, Long length, List<ByteBuffer> buffers) {
+    private ByteBuffersAsyncRequestBody(String mimetype,
+                                        Long length,
+                                        List<ByteBuffer> buffers) {
         this.mimetype = mimetype;
         this.buffers = buffers;
         this.length = length;
@@ -121,6 +123,10 @@ public String body() {
         return BodyType.BYTES.getName();
     }
 
+    public static ByteBuffersAsyncRequestBody of(List<ByteBuffer> buffers, long length) {
+        return new ByteBuffersAsyncRequestBody(Mimetype.MIMETYPE_OCTET_STREAM, length, buffers);
+    }
+
     public static ByteBuffersAsyncRequestBody of(List<ByteBuffer> buffers) {
         long length = buffers.stream()
                             .mapToLong(ByteBuffer::remaining)
@@ -129,7 +135,11 @@ public static ByteBuffersAsyncRequestBody of(List<ByteBuffer> buffers) {
     }
 
     public static ByteBuffersAsyncRequestBody of(ByteBuffer... buffers) {
-        return of(Arrays.asList(buffers));
+        List<ByteBuffer> byteBuffers = Arrays.asList(buffers);
+        long length = byteBuffers.stream()
+                             .mapToLong(ByteBuffer::remaining)
+                             .sum();
+        return of(byteBuffers, length);
     }
 
     public static ByteBuffersAsyncRequestBody of(Long length, ByteBuffer... buffers) {

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

@@ -34,6 +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.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.internal.util.Mimetype;
 import software.amazon.awssdk.core.internal.util.NoopSubscription;
@@ -86,6 +87,11 @@ public SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration s
         return new FileAsyncRequestBodySplitHelper(this, splitConfiguration).split();
     }
 
+    @Override
+    public SdkPublisher<CloseableAsyncRequestBody> splitCloseable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+        return split(splitConfiguration).map(body -> new ClosableAsyncRequestBodyWrapper(body));
+    }
+
     public Path path() {
         return path;
     }
@@ -436,4 +442,32 @@ private void signalOnError(Throwable t) {
     private static AsynchronousFileChannel openInputChannel(Path path) throws IOException {
         return AsynchronousFileChannel.open(path, StandardOpenOption.READ);
     }
+
+    private static class ClosableAsyncRequestBodyWrapper implements CloseableAsyncRequestBody {
+        private final AsyncRequestBody delegate;
+
+        ClosableAsyncRequestBodyWrapper(AsyncRequestBody body) {
+            this.delegate = body;
+        }
+
+        @Override
+        public Optional<Long> contentLength() {
+            return delegate.contentLength();
+        }
+
+        @Override
+        public void subscribe(Subscriber<? super ByteBuffer> s) {
+            delegate.subscribe(s);
+        }
+
+        @Override
+        public void close() {
+            // no op
+        }
+
+        @Override
+        public String body() {
+            return delegate.body();
+        }
+    }
 }