Fix build and add more tests

Author: zoewangg

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.ClosableAsyncRequestBody;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
@@ -76,6 +77,16 @@ public SdkPublisher<AsyncRequestBody> split(Consumer<AsyncRequestBodySplitConfig
             return delegate.split(splitConfiguration);
         }
 
+        @Override
+        public SdkPublisher<ClosableAsyncRequestBody> splitV2(AsyncRequestBodySplitConfiguration splitConfiguration) {
+            return delegate.splitV2(splitConfiguration);
+        }
+
+        @Override
+        public SdkPublisher<ClosableAsyncRequestBody> splitV2(Consumer<AsyncRequestBodySplitConfiguration.Builder> splitConfiguration) {
+            return delegate.splitV2(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/SplittingPublisher.java

@@ -65,7 +65,7 @@ public SplittingPublisher(AsyncRequestBody asyncRequestBody,
 
         Validate.isTrue(bufferSizeInBytes >= chunkSizeInBytes,
                         "bufferSizeInBytes must be larger than or equal to " +
-                        "chunkSizeInBytes if the content length is unknown");
+                        "chunkSizeInBytes");
     }
 
     @Override
@@ -77,7 +77,7 @@ public void subscribe(Subscriber<? super ClosableAsyncRequestBody> downstreamSub
     private class SplittingSubscriber implements Subscriber<ByteBuffer> {
         private Subscription upstreamSubscription;
         private final Long upstreamSize;
-        private final AtomicInteger chunkNumber = new AtomicInteger(0);
+        private final AtomicInteger chunkNumber = new AtomicInteger(1);
         private volatile DownstreamBody currentBody;
         private final AtomicBoolean hasOpenUpstreamDemand = new AtomicBoolean(false);
         private final AtomicLong dataBuffered = new AtomicLong(0);
@@ -155,7 +155,7 @@ public void onNext(ByteBuffer byteBuffer) {
 
         private void completeCurrentBodyAndCreateNewIfNeeded(ByteBuffer byteBuffer) {
             completeCurrentBodyAndDeliver();
-            int currentChunk = chunkNumber.incrementAndGet();
+            int nextChunk = chunkNumber.incrementAndGet();
             boolean shouldCreateNewDownstreamRequestBody;
             Long dataRemaining = totalDataRemaining();
 
@@ -167,7 +167,7 @@ private void completeCurrentBodyAndCreateNewIfNeeded(ByteBuffer byteBuffer) {
 
             if (shouldCreateNewDownstreamRequestBody) {
                 long chunkSize = calculateChunkSize(dataRemaining);
-                currentBody = initializeNextDownstreamBody(upstreamSize != null, chunkSize, currentChunk);
+                currentBody = initializeNextDownstreamBody(upstreamSize != null, chunkSize, nextChunk);
             }
         }
 

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/ClosableAsyncRequestBodyAdaptorTest.java

@@ -20,26 +20,43 @@
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doNothing;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import io.reactivex.Flowable;
 import io.reactivex.FlowableSubscriber;
+import io.reactivex.internal.observers.BiConsumerSingleObserver;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
+import java.util.Observable;
+import java.util.Observer;
+import java.util.Optional;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
 import software.amazon.awssdk.core.async.ClosableAsyncRequestBody;
 import software.amazon.awssdk.core.exception.NonRetryableException;
 
 public class ClosableAsyncRequestBodyAdaptorTest {
+    private ClosableAsyncRequestBody closableAsyncRequestBody;
 
-    @Test
-    void resubscribe_shouldThrowException() {
-        ClosableAsyncRequestBody closableAsyncRequestBody = Mockito.mock(ClosableAsyncRequestBody.class);
+    @BeforeEach
+    public void setup() {
+        closableAsyncRequestBody =Mockito.mock(ClosableAsyncRequestBody.class);
         Mockito.when(closableAsyncRequestBody.doAfterOnComplete(any(Runnable.class))).thenReturn(closableAsyncRequestBody);
         Mockito.when(closableAsyncRequestBody.doAfterOnCancel(any(Runnable.class))).thenReturn(closableAsyncRequestBody);
         Mockito.when(closableAsyncRequestBody.doAfterOnError(any(Consumer.class))).thenReturn(closableAsyncRequestBody);
+    }
 
+    @Test
+    void resubscribe_shouldThrowException() {
         ClosableAsyncRequestBodyAdaptor adaptor = new ClosableAsyncRequestBodyAdaptor(closableAsyncRequestBody);
         Subscriber subscriber = Mockito.mock(Subscriber.class);
         adaptor.subscribe(subscriber);
@@ -55,4 +72,99 @@ void resubscribe_shouldThrowException() {
             .hasMessageContaining("A retry was attempted");
     }
 
+    @Test
+    void onComplete_shouldCloseAsyncRequestBody() {
+        TestClosableAsyncRequestBody asyncRequestBody = new TestClosableAsyncRequestBody();
+        ClosableAsyncRequestBodyAdaptor adaptor = new ClosableAsyncRequestBodyAdaptor(asyncRequestBody);
+        CompletableFuture<byte[]> future = new CompletableFuture<>();
+        Subscriber<ByteBuffer> subscriber = new ByteArrayAsyncResponseTransformer.BaosSubscriber(future);
+        adaptor.subscribe(subscriber);
+        assertThat(asyncRequestBody.closeInvoked).isTrue();
+    }
+
+    @Test
+    void cancel_shouldCloseAsyncRequestBody() {
+        TestClosableAsyncRequestBody asyncRequestBody = new TestClosableAsyncRequestBody();
+        ClosableAsyncRequestBodyAdaptor adaptor = new ClosableAsyncRequestBodyAdaptor(asyncRequestBody);
+        Subscriber<ByteBuffer> subscriber = new Subscriber<ByteBuffer>() {
+            @Override
+            public void onSubscribe(Subscription s) {
+                s.cancel();
+            }
+
+            @Override
+            public void onNext(ByteBuffer byteBuffer) {
+            }
+
+            @Override
+            public void onError(Throwable t) {
+            }
+
+            @Override
+            public void onComplete() {
+            }
+        };
+        adaptor.subscribe(subscriber);
+        assertThat(asyncRequestBody.closeInvoked).isTrue();
+    }
+
+    @Test
+    void onError_shouldCloseAsyncRequestBody() {
+        OnErrorClosableAsyncRequestBody asyncRequestBody = new OnErrorClosableAsyncRequestBody();
+        ClosableAsyncRequestBodyAdaptor adaptor = new ClosableAsyncRequestBodyAdaptor(asyncRequestBody);
+        CompletableFuture<byte[]> future = new CompletableFuture<>();
+        Subscriber<ByteBuffer> subscriber = new ByteArrayAsyncResponseTransformer.BaosSubscriber(future);
+        adaptor.subscribe(subscriber);
+        assertThat(asyncRequestBody.closeInvoked).isTrue();
+    }
+
+
+    private static class TestClosableAsyncRequestBody implements ClosableAsyncRequestBody {
+        private boolean closeInvoked;
+
+        @Override
+        public Optional<Long> contentLength() {
+            return Optional.empty();
+        }
+
+        @Override
+        public void subscribe(Subscriber<? super ByteBuffer> s) {
+            Flowable.just(ByteBuffer.wrap("foo bar".getBytes(StandardCharsets.UTF_8)))
+                    .subscribe(s);
+        }
+
+        @Override
+        public void close() {
+            closeInvoked = true;
+        }
+    }
+
+    private static class OnErrorClosableAsyncRequestBody implements ClosableAsyncRequestBody {
+        private boolean closeInvoked;
+
+        @Override
+        public Optional<Long> contentLength() {
+            return Optional.empty();
+        }
+
+        @Override
+        public void subscribe(Subscriber<? super ByteBuffer> s) {
+            s.onSubscribe(new Subscription() {
+                @Override
+                public void request(long n) {
+                    s.onError(new IllegalStateException("foobar"));
+                }
+
+                @Override
+                public void cancel() {
+
+                }
+            });
+        }
+
+        @Override
+        public void close() {
+            closeInvoked = true;
+        }
+    }
 }

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java

@@ -151,6 +151,10 @@ public void onNext(ClosableAsyncRequestBody asyncRequestBody) {
         }
 
         int currentPartNum = partNumber.getAndIncrement();
+
+        log.debug(() -> String.format("Received asyncRequestBody for part number %d with length %s", currentPartNum,
+                                      asyncRequestBody.contentLength()));
+
         if (existingParts.containsKey(currentPartNum)) {
             asyncRequestBody.subscribe(new CancelledSubscriber<>());
             asyncRequestBody.contentLength().ifPresent(progressListener::subscriberOnNext);

services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UploadWithUnknownContentLengthHelper.java

@@ -159,7 +159,7 @@ public void onNext(ClosableAsyncRequestBody asyncRequestBody) {
                 return;
             }
             int currentPartNum = partNumber.incrementAndGet();
-            log.debug(() -> String.format("Received asyncRequestBody for part number %d with length %d", currentPartNum,
+            log.debug(() -> String.format("Received asyncRequestBody for part number %d with length %s", currentPartNum,
                       asyncRequestBody.contentLength()));
             asyncRequestBodyInFlight.incrementAndGet();