Fix and clean up

Author: zoewangg

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

@@ -30,7 +30,6 @@
 import software.amazon.awssdk.core.internal.async.ByteArrayAsyncRequestBody;
 import software.amazon.awssdk.core.internal.async.FileAsyncRequestBody;
 import software.amazon.awssdk.core.internal.async.InputStreamWithExecutorAsyncRequestBody;
-import software.amazon.awssdk.core.internal.async.SplittingPublisher;
 import software.amazon.awssdk.core.internal.util.Mimetype;
 import software.amazon.awssdk.utils.BinaryUtils;
 
@@ -70,11 +69,6 @@ default String contentType() {
         return Mimetype.MIMETYPE_OCTET_STREAM;
     }
 
-    default SdkPublisher<AsyncRequestBody> split(long partSizeInBytes,
-                                                 long maxMemoryUsageInBytes) {
-        return new SplittingPublisher(this, partSizeInBytes, maxMemoryUsageInBytes);
-    }
-
     /**
      * Creates an {@link AsyncRequestBody} the produces data from the input ByteBuffer publisher.
      * The data is delivered when the publisher publishes the data.

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

@@ -17,27 +17,55 @@
 
 import java.nio.ByteBuffer;
 import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
+import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
+import software.amazon.awssdk.utils.Logger;
+import software.amazon.awssdk.utils.Validate;
 import software.amazon.awssdk.utils.async.SimplePublisher;
 
+/**
+ * Splits an {@link SdkPublisher} to multiple smaller {@link AsyncRequestBody}s, each of which publishes a specific portion of the
+ * original data.
+ * // TODO: create a static factory method in AsyncRequestBody for this
+ * // TODO: fix the case where content length is null
+ */
+@SdkInternalApi
 public class SplittingPublisher implements SdkPublisher<AsyncRequestBody> {
+    private static final Logger log = Logger.loggerFor(SplittingPublisher.class);
     private final AsyncRequestBody upstreamPublisher;
-    private final SplittingSubscriber splittingSubscriber = new SplittingSubscriber();
+    private final SplittingSubscriber splittingSubscriber;
     private final SimplePublisher<AsyncRequestBody> downstreamPublisher = new SimplePublisher<>();
-    private final long partSizeInBytes;
+    private final long chunkSizeInBytes;
     private final long maxMemoryUsageInBytes;
+    private final CompletableFuture<Void> future;
+
+    private SplittingPublisher(Builder builder) {
+        this.upstreamPublisher =  Validate.paramNotNull(builder.asyncRequestBody, "asyncRequestBody");
+        this.chunkSizeInBytes = Validate.paramNotNull(builder.chunkSizeInBytes, "chunkSizeInBytes");
+        this.splittingSubscriber = new SplittingSubscriber(upstreamPublisher.contentLength().orElse(null));
+        this.maxMemoryUsageInBytes = builder.maxMemoryUsageInBytes == null ? Long.MAX_VALUE : builder.maxMemoryUsageInBytes;
+        this.future = builder.future;
+
+        // We need to cancel upstream subscription if the future gets cancelled.
+        future.whenComplete((r, t) -> {
+            if (t != null) {
+                if (splittingSubscriber.upstreamSubscription != null) {
+                    log.trace(() -> "Cancelling subscription because return future completed exceptionally ", t);
+                    splittingSubscriber.upstreamSubscription.cancel();
+                }
+            }
+        });
+    }
 
-    public SplittingPublisher(AsyncRequestBody asyncRequestBody,
-                              long partSizeInBytes,
-                              long maxMemoryUsageInBytes) {
-        this.upstreamPublisher = asyncRequestBody;
-        this.partSizeInBytes = partSizeInBytes;
-        this.maxMemoryUsageInBytes = maxMemoryUsageInBytes;
+    public static Builder builder() {
+        return new Builder();
     }
 
     @Override
@@ -48,50 +76,70 @@ public void subscribe(Subscriber<? super AsyncRequestBody> downstreamSubscriber)
 
     private class SplittingSubscriber implements Subscriber<ByteBuffer> {
         private Subscription upstreamSubscription;
-        private final Long upstreamSize = upstreamPublisher.contentLength().orElse(null);
-
-        private int partNumber = 0;
-        private DownstreamBody currentBody;
-
+        private final Long upstreamSize;
+        private final AtomicInteger chunkNumber = new AtomicInteger(0);
+        private volatile DownstreamBody currentBody;
         private final AtomicBoolean hasOpenUpstreamDemand = new AtomicBoolean(false);
         private final AtomicLong dataBuffered = new AtomicLong(0);
 
+        /**
+         * A hint to determine whether we will exceed maxMemoryUsage by the next OnNext call.
+         */
+        private int byteBufferSizeHint;
+
+        SplittingSubscriber(Long upstreamSize) {
+            this.upstreamSize = upstreamSize;
+        }
+
         @Override
         public void onSubscribe(Subscription s) {
             this.upstreamSubscription = s;
-            this.currentBody = new DownstreamBody(calculatePartSize());
+            this.currentBody = new DownstreamBody(calculateChunkSize(), chunkNumber.get());
+            sendCurrentBody();
+            // We need to request subscription *after* we set currentBody because onNext could be invoked right away.
+            upstreamSubscription.request(1);
         }
 
         @Override
         public void onNext(ByteBuffer byteBuffer) {
             hasOpenUpstreamDemand.set(false);
+            byteBufferSizeHint = byteBuffer.remaining();
 
             while (true) {
-                int amountRemainingInPart = Math.toIntExact(partSizeInBytes - currentBody.partLength);
+                int amountRemainingInPart = amountRemainingInPart();
+                int finalAmountRemainingInPart = amountRemainingInPart;
+                if (amountRemainingInPart == 0) {
+                    currentBody.complete();
+                    int currentChunk = chunkNumber.incrementAndGet();
+                    Long partSize = calculateChunkSize();
+                    currentBody = new DownstreamBody(partSize, currentChunk);
+                    sendCurrentBody();
+                }
 
-                if (amountRemainingInPart < byteBuffer.remaining()) {
-                    // TODO: should we avoid sending empty byte buffers, which can happen here?
-                    currentBody.send(byteBuffer);
+                amountRemainingInPart = amountRemainingInPart();
+                if (amountRemainingInPart >= byteBuffer.remaining()) {
+                    currentBody.send(byteBuffer.duplicate());
                     break;
                 }
 
                 ByteBuffer firstHalf = byteBuffer.duplicate();
                 int newLimit = firstHalf.position() + amountRemainingInPart;
-                firstHalf.limit(newLimit); // TODO: Am I off by one here?
+                firstHalf.limit(newLimit);
+                byteBuffer.position(newLimit);
                 currentBody.send(firstHalf);
-                currentBody.complete();
-
-                ++partNumber;
-                currentBody = new DownstreamBody(calculatePartSize());
-                downstreamPublisher.send(currentBody);
-                byteBuffer.position(newLimit); // TODO: Am I off by one here?
             }
 
             maybeRequestMoreUpstreamData();
         }
 
+        private int amountRemainingInPart() {
+            return Math.toIntExact(currentBody.totalLength - currentBody.transferredLength);
+        }
+
         @Override
         public void onComplete() {
+            log.trace(() -> "Received onComplete()");
+            downstreamPublisher.complete().thenRun(() -> future.complete(null));
             currentBody.complete();
         }
 
@@ -100,53 +148,73 @@ public void onError(Throwable t) {
             currentBody.error(t);
         }
 
-        private Long calculatePartSize() {
+        private void sendCurrentBody() {
+            downstreamPublisher.send(currentBody).exceptionally(t -> {
+                downstreamPublisher.error(t);
+                return null;
+            });
+        }
+
+        private Long calculateChunkSize() {
             Long dataRemaining = dataRemaining();
             if (dataRemaining == null) {
                 return null;
             }
 
-            return Math.min(partSizeInBytes, dataRemaining);
+            return Math.min(chunkSizeInBytes, dataRemaining);
         }
 
         private void maybeRequestMoreUpstreamData() {
-            if (dataBuffered.get() < maxMemoryUsageInBytes && hasOpenUpstreamDemand.compareAndSet(false, true)) {
-                // TODO: max memory usage might not be the best name, since we can technically go a little above
-                // this limit when we add on a new byte buffer. But we don't know what the size of a buffer we request
-                // will be, so I don't think we can have a truly accurate max. Maybe we call it minimum buffer size instead?
+            long buffered = dataBuffered.get();
+            if (shouldRequestMoreData(buffered) &&
+                hasOpenUpstreamDemand.compareAndSet(false, true)) {
+                log.trace(() -> "Requesting more data, current data buffered: " + buffered);
                 upstreamSubscription.request(1);
             }
         }
 
+        private boolean shouldRequestMoreData(long buffered) {
+            return buffered == 0 || buffered + byteBufferSizeHint < maxMemoryUsageInBytes;
+        }
+
         private Long dataRemaining() {
             if (upstreamSize == null) {
                 return null;
             }
-            return upstreamSize - (partNumber * partSizeInBytes);
+            return upstreamSize - (chunkNumber.get() * chunkSizeInBytes);
         }
 
         private class DownstreamBody implements AsyncRequestBody {
-            private final Long size;
+            private final Long totalLength;
             private final SimplePublisher<ByteBuffer> delegate = new SimplePublisher<>();
-            private long partLength = 0;
+            private final int chunkNumber;
+            private volatile long transferredLength = 0;
 
-            private DownstreamBody(Long size) {
-                this.size = size;
+            private DownstreamBody(Long totalLength, int chunkNumber) {
+                this.totalLength = totalLength;
+                this.chunkNumber = chunkNumber;
             }
 
             @Override
             public Optional<Long> contentLength() {
-                return Optional.ofNullable(size);
+                return Optional.ofNullable(totalLength);
             }
 
             public void send(ByteBuffer data) {
+                log.trace(() -> "Sending bytebuffer " + data);
                 int length = data.remaining();
-                partLength += length;
+                transferredLength += length;
                 addDataBuffered(length);
-                delegate.send(data).thenRun(() -> addDataBuffered(-length));
+                delegate.send(data).whenComplete((r, t) -> {
+                    addDataBuffered(-length);
+                    if (t != null) {
+                        error(t);
+                    }
+                });
             }
 
             public void complete() {
+                log.debug(() -> "Received complete() for chunk number: " + chunkNumber);
                 delegate.complete();
             }
 
@@ -167,4 +235,64 @@ private void addDataBuffered(int length) {
             }
         }
     }
+    
+    public static final class Builder {
+        private AsyncRequestBody asyncRequestBody;
+        private Long chunkSizeInBytes;
+        private Long maxMemoryUsageInBytes;
+        private CompletableFuture<Void> future;
+
+        /**
+         * Configures the asyncRequestBody to split
+         *
+         * @param asyncRequestBody The new asyncRequestBody value.
+         * @return This object for method chaining.
+         */
+        public Builder asyncRequestBody(AsyncRequestBody asyncRequestBody) {
+            this.asyncRequestBody = asyncRequestBody;
+            return this;
+        }
+
+        /**
+         * Configures the size of the chunk for each {@link AsyncRequestBody} to publish
+         *
+         * @param chunkSizeInBytes The new chunkSizeInBytes value.
+         * @return This object for method chaining.
+         */
+        public Builder chunkSizeInBytes(Long chunkSizeInBytes) {
+            this.chunkSizeInBytes = chunkSizeInBytes;
+            return this;
+        }
+
+        /**
+         * Sets the maximum memory usage in bytes. By default, it uses unlimited memory.
+         *
+         * @param maxMemoryUsageInBytes The new maxMemoryUsageInBytes value.
+         * @return This object for method chaining.
+         */
+        // TODO: max memory usage might not be the best name, since we may technically go a little above this limit when we add
+        //  on a new byte buffer. But we don't know for sure what the size of a buffer we request will be (we do use the size
+        //  for the last byte buffer as a hint), so I don't think we can have a truly accurate max. Maybe we call it minimum
+        //  buffer size instead?
+        public Builder maxMemoryUsageInBytes(Long maxMemoryUsageInBytes) {
+            this.maxMemoryUsageInBytes = maxMemoryUsageInBytes;
+            return this;
+        }
+
+        /**
+         * Sets the result future. The future will be completed when all request bodies
+         * have been sent.
+         *
+         * @param future The new future value.
+         * @return This object for method chaining.
+         */
+        public Builder resultFuture(CompletableFuture<Void> future) {
+            this.future = future;
+            return this;
+        }
+
+        public SplittingPublisher build() {
+            return new SplittingPublisher(this);
+        }
+    }
 }