Refactoring based on API surface are review

Author: zoewangg

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

@@ -33,7 +33,6 @@
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.core.FileRequestBodyConfiguration;
 import software.amazon.awssdk.core.internal.async.ByteBuffersAsyncRequestBody;
-import software.amazon.awssdk.core.internal.async.ClosableAsyncRequestBodyAdaptor;
 import software.amazon.awssdk.core.internal.async.FileAsyncRequestBody;
 import software.amazon.awssdk.core.internal.async.InputStreamWithExecutorAsyncRequestBody;
 import software.amazon.awssdk.core.internal.async.SplittingPublisher;
@@ -501,22 +500,23 @@ 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
      * is 2MB and the default buffer size is 8MB.
      *
      * <p>
-     * Each divided {@link AsyncRequestBody} is sent after the entire content for that chunk is buffered.
+     * By default, if content length of this {@link AsyncRequestBody} is present, each divided {@link AsyncRequestBody} is
+     * delivered to the subscriber right after it's initialized. On the other hand, if content length is null, it is sent after
+     * the entire content for that chunk is buffered. In this case, the configured {@code maxMemoryUsageInBytes} must be larger
+     * 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 #splitV2(AsyncRequestBodySplitConfiguration)} instead.
      */
-    @Deprecated
     default SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
-        return splitV2(splitConfiguration).map(body -> new ClosableAsyncRequestBodyAdaptor(body));
+        return new SplittingPublisher(this, splitConfiguration, false).map(r -> r);
     }
 
     /**
@@ -525,27 +525,26 @@ default SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration
      * size is 2MB and the default buffer size is 8MB.
      *
      * <p>
-     * Each divided {@link ClosableAsyncRequestBody} is sent after the entire content for that chunk is buffered. This behavior
-     * may be different if a specific implementation of this interface overrides this method.
+     * The default implementation behaves the same as {@link #split(AsyncRequestBodySplitConfiguration)}. This behavior may
+     * vary in different implementations.
      *
      * <p>
-     * Each {@link ClosableAsyncRequestBody} MUST be closed by the user when it is ready to be disposed.
+     * Caller is responsible for closing {@link ClosableAsyncRequestBody} when it is ready to be disposed to release any
+     * resources.
      *
      * @see AsyncRequestBodySplitConfiguration
      */
-    default SdkPublisher<ClosableAsyncRequestBody> splitV2(AsyncRequestBodySplitConfiguration splitConfiguration) {
+    default SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
-        return new SplittingPublisher(this, splitConfiguration);
+        return new SplittingPublisher(this, splitConfiguration, false);
     }
 
     /**
      * 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 #splitV2(Consumer)} instead.
      */
-    @Deprecated
     default SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
         return split(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
@@ -555,12 +554,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 #splitV2(Consumer)
+     * @see #splitClosable(Consumer)
      */
-    default SdkPublisher<ClosableAsyncRequestBody> splitV2(
+    default SdkPublisher<ClosableAsyncRequestBody> splitClosable(
         Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
         Validate.notNull(splitConfiguration, "splitConfiguration");
-        return splitV2(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
+        return splitClosable(AsyncRequestBodySplitConfiguration.builder().applyMutation(splitConfiguration).build());
     }
 
     @SdkProtectedApi
@@ -569,6 +568,8 @@ enum BodyType {
         BYTES("Bytes", "b"),
         STREAM("Stream", "s"),
         PUBLISHER("Publisher", "p"),
+        RETRYABLE("RetryableSub", "rs"),
+        NON_RETRYABLE("NonRetryableSub", "nrs"),
         UNKNOWN("Unknown", "u");
 
         private static final Map<String, BodyType> VALUE_MAP =

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

@@ -0,0 +1,59 @@
+/*
+ * 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;
+
+/**
+ * An {@link AsyncRequestBody} decorator that can be split into buffered sub {@link AsyncRequestBody}s. Each sub
+ * {@link AsyncRequestBody} can be retried/resubscribed.
+ */
+@SdkPublicApi
+public final class BufferedSplittableAsyncRequestBody implements AsyncRequestBody {
+    private final AsyncRequestBody delegate;
+
+    private BufferedSplittableAsyncRequestBody(AsyncRequestBody delegate) {
+        this.delegate = delegate;
+    }
+
+    public static BufferedSplittableAsyncRequestBody create(AsyncRequestBody delegate) {
+        return new BufferedSplittableAsyncRequestBody(delegate);
+    }
+
+    @Override
+    public Optional<Long> contentLength() {
+        return delegate.contentLength();
+    }
+
+    @Override
+    public SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+        return new SplittingPublisher(this, splitConfiguration, true);
+    }
+
+    @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/listener/AsyncRequestBodyListener.java

@@ -78,14 +78,14 @@ public SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfig
         }
 
         @Override
-        public SdkPublisher<ClosableAsyncRequestBody> splitV2(AsyncRequestBodySplitConfiguration splitConfiguration) {
-            return delegate.splitV2(splitConfiguration);
+        public SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
+            return delegate.splitClosable(splitConfiguration);
         }
 
         @Override
-        public SdkPublisher<ClosableAsyncRequestBody> splitV2(
+        public SdkPublisher<ClosableAsyncRequestBody> splitClosable(
             Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
-            return delegate.splitV2(splitConfiguration);
+            return delegate.splitClosable(splitConfiguration);
         }
 
         @Override

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

@@ -123,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)

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

@@ -88,7 +88,7 @@ public SdkPublisher<AsyncRequestBody> split(AsyncRequestBodySplitConfiguration s
     }
 
     @Override
-    public SdkPublisher<ClosableAsyncRequestBody> splitV2(AsyncRequestBodySplitConfiguration splitConfiguration) {
+    public SdkPublisher<ClosableAsyncRequestBody> splitClosable(AsyncRequestBodySplitConfiguration splitConfiguration) {
         return split(splitConfiguration).map(body -> new ClosableAsyncRequestBodyWrapper(body));
     }
 

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

@@ -0,0 +1,128 @@
+/*
+ * 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.internal.async;
+
+import java.nio.ByteBuffer;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Consumer;
+import org.reactivestreams.Subscriber;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.exception.NonRetryableException;
+import software.amazon.awssdk.core.internal.util.NoopSubscription;
+import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.async.SimplePublisher;
+
+/**
+ * A {@link SubAsyncRequestBody} implementation that doesn't resubscribe/retry
+ */
+@SdkInternalApi
+public final class NonRetryableSubAsyncRequestBody implements SubAsyncRequestBody {
+    private static final Logger log = Logger.loggerFor(NonRetryableSubAsyncRequestBody.class);
+    private final long maxLength;
+    private final Long totalLength;
+    private final SimplePublisher<ByteBuffer> delegate = new SimplePublisher<>();
+    private final int partNumber;
+    private final AtomicBoolean subscribeCalled = new AtomicBoolean(false);
+    private final Consumer<Long> onNumBytesReceived;
+    private final Consumer<Long> onNumBytesConsumed;
+    private final boolean contentLengthKnown;
+    private final String sourceBodyName;
+    private volatile long bufferedLength = 0;
+
+    // TODO: builder pattern
+    public NonRetryableSubAsyncRequestBody(boolean contentLengthKnown,
+                                           long maxLength,
+                                           int partNumber,
+                                           Consumer<Long> onNumBytesReceived,
+                                           Consumer<Long> onNumBytesConsumed, String sourceBodyName) {
+        this.contentLengthKnown = contentLengthKnown;
+        this.totalLength = contentLengthKnown ? maxLength : null;
+        this.maxLength = maxLength;
+        this.partNumber = partNumber;
+        this.onNumBytesReceived = onNumBytesReceived;
+        this.onNumBytesConsumed = onNumBytesConsumed;
+        this.sourceBodyName = sourceBodyName;
+    }
+
+    @Override
+    public Optional<Long> contentLength() {
+        return totalLength != null ? Optional.of(totalLength) : Optional.of(bufferedLength);
+    }
+
+    public void send(ByteBuffer data) {
+        log.debug(() -> String.format("Sending bytebuffer %s to part %d", data, partNumber));
+        int length = data.remaining();
+        bufferedLength += length;
+        onNumBytesReceived.accept((long) length);
+        delegate.send(data).whenComplete((r, t) -> {
+            onNumBytesConsumed.accept((long) length);
+            if (t != null) {
+                error(t);
+            }
+        });
+    }
+
+    public void complete() {
+        log.debug(() -> "Received complete() for part number: " + partNumber + " length " + bufferedLength);
+        delegate.complete().whenComplete((r, t) -> {
+            if (t != null) {
+                error(t);
+            }
+        });
+    }
+
+    @Override
+    public long maxLength() {
+        return maxLength;
+    }
+
+    @Override
+    public long receivedBytesLength() {
+        return bufferedLength;
+    }
+
+    @Override
+    public boolean contentLengthKnown() {
+        return contentLengthKnown;
+    }
+
+    @Override
+    public int partNumber() {
+        return partNumber;
+    }
+
+    public void error(Throwable error) {
+        delegate.error(error);
+    }
+
+    @Override
+    public void subscribe(Subscriber<? super ByteBuffer> s) {
+        if (subscribeCalled.compareAndSet(false, true)) {
+            delegate.subscribe(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. "));
+        }
+    }
+
+    @Override
+    public String body() {
+        return sourceBodyName;
+    }
+}