Implementation of MultipartDownloaderSubscriber (#4931)

* Implementation of MultipartDownloaderSubscriber and previous SplittingTransformer PR comments. Wiremock tests.
- Use builder and template methods in DelegatingBufferingSubscriber

* fix DelegatingBufferingSubscriber not considering downstream demand

* fix DelegatingBufferingSubscriber not always respecting downstream demand

* test for missing byte bug

* test for missing byte bug

* fix merge

* remove buffering to see if byte length error is still triggered

* missing byte investigation

* testing errors with InputStream and Publisher

* checkstyle

* checkstyle

* prevent test memory heap error

* fix merge

* handle cancel signal correctly, and make IndividualPartSubscriber non-static to do so.
- removed trace/debug logging

* attempt fix thread hanging for PublisherAsyncResponseTransformer

Author: L-Applin

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

@@ -83,7 +83,7 @@ public class SplittingTransformer<ResponseT, ResultT> implements SdkPublisher<As
     /**
      * The buffer size used to buffer the content received from the downstream subscriber
      */
-    private final long maximumBufferSize;
+    private final long maximumBufferInBytes;
 
     /**
      * This publisher is used to send the bytes received from the downstream subscriber's transformers to a
@@ -116,7 +116,7 @@ private SplittingTransformer(AsyncResponseTransformer<ResponseT, ResultT> upstre
         this.resultFuture = Validate.paramNotNull(
             resultFuture, "resultFuture");
         Validate.notNull(maximumBufferSizeInBytes, "maximumBufferSizeInBytes");
-        this.maximumBufferSize = Validate.isPositive(
+        this.maximumBufferInBytes = Validate.isPositive(
             maximumBufferSizeInBytes, "maximumBufferSizeInBytes");
     }
 
@@ -160,8 +160,7 @@ public void request(long n) {
         public void cancel() {
             if (isCancelled.compareAndSet(false, true)) {
                 log.trace(() -> "Cancelling splitting transformer");
-                publisherToUpstream.complete();
-                downstreamSubscriber = null;
+                handleCancelState();
             }
         }
     }
@@ -196,25 +195,53 @@ private boolean doEmit() {
         return false;
     }
 
+    private void handleCancelState() {
+        synchronized (this) {
+            if (downstreamSubscriber == null) {
+                return;
+            }
+            if (!onStreamCalled.get()) {
+                // we never subscribe publisherToUpstream to the upstream, it would not complete
+                downstreamSubscriber = null;
+                return;
+            }
+            publisherToUpstream.complete().whenComplete((v, t) -> {
+                if (downstreamSubscriber == null) {
+                    return;
+                }
+                if (t != null) {
+                    downstreamSubscriber.onError(t);
+                } else {
+                    downstreamSubscriber.onComplete();
+                }
+                downstreamSubscriber = null;
+            });
+        }
+    }
+
     /**
      * The AsyncResponseTransformer for each of the individual requests that is sent back to the downstreamSubscriber when
      * requested. A future is created per request that is completed when onComplete is called on the subscriber for that request
      * body publisher.
      */
     private class IndividualTransformer implements AsyncResponseTransformer<ResponseT, ResponseT> {
-
         private ResponseT response;
         private CompletableFuture<ResponseT> individualFuture;
 
         @Override
         public CompletableFuture<ResponseT> prepare() {
             this.individualFuture = new CompletableFuture<>();
             if (preparedCalled.compareAndSet(false, true)) {
-                log.trace(() -> "calling prepare on the upstream transformer");
                 CompletableFuture<ResultT> upstreamFuture = upstreamResponseTransformer.prepare();
-                CompletableFutureUtils.forwardExceptionTo(resultFuture, upstreamFuture);
-                CompletableFutureUtils.forwardResultTo(upstreamFuture, resultFuture);
+                if (!resultFuture.isDone()) {
+                    CompletableFutureUtils.forwardResultTo(upstreamFuture, resultFuture);
+                }
             }
+            individualFuture.whenComplete((r, e) -> {
+                if (isCancelled.get()) {
+                    handleCancelState();
+                }
+            });
             return this.individualFuture;
         }
 
@@ -229,38 +256,40 @@ public void onResponse(ResponseT response) {
 
         @Override
         public void onStream(SdkPublisher<ByteBuffer> publisher) {
+            if (downstreamSubscriber == null) {
+                return;
+            }
             if (onStreamCalled.compareAndSet(false, true)) {
                 log.trace(() -> "calling onStream on the upstream transformer");
-                upstreamResponseTransformer.onStream(
-                    upstreamSubscriber ->
-                        publisherToUpstream.subscribe(new DelegatingBufferingSubscriber(maximumBufferSize, upstreamSubscriber))
-                );
+                upstreamResponseTransformer.onStream(upstreamSubscriber -> publisherToUpstream.subscribe(
+                    DelegatingBufferingSubscriber.builder()
+                                                 .maximumBufferInBytes(maximumBufferInBytes)
+                                                 .delegate(upstreamSubscriber)
+                                                 .build()
+                ));
             }
-            publisher.subscribe(new IndividualPartSubscriber<>(this.individualFuture, response, publisherToUpstream));
+            publisher.subscribe(new IndividualPartSubscriber<>(this.individualFuture, response));
         }
 
         @Override
         public void exceptionOccurred(Throwable error) {
-            individualFuture.completeExceptionally(error);
+            publisherToUpstream.error(error);
             upstreamResponseTransformer.exceptionOccurred(error);
         }
     }
 
     /**
      * the Subscriber for each of the individual request's ByteBuffer publisher
      */
-    static class IndividualPartSubscriber<T> implements Subscriber<ByteBuffer> {
+    class IndividualPartSubscriber<T> implements Subscriber<ByteBuffer> {
 
         private final CompletableFuture<T> future;
         private final T response;
-        private final SimplePublisher<ByteBuffer> bodyPartPublisher;
         private Subscription subscription;
 
-        IndividualPartSubscriber(CompletableFuture<T> future, T response,
-                                 SimplePublisher<ByteBuffer> bodyPartPublisher) {
+        IndividualPartSubscriber(CompletableFuture<T> future, T response) {
             this.future = future;
             this.response = response;
-            this.bodyPartPublisher = bodyPartPublisher;
         }
 
         @Override
@@ -270,16 +299,15 @@ public void onSubscribe(Subscription s) {
                 return;
             }
             this.subscription = s;
-            // request everything, data will be buffered by the DelegatingBufferingSubscriber
-            s.request(Long.MAX_VALUE);
+            s.request(1);
         }
 
         @Override
         public void onNext(ByteBuffer byteBuffer) {
             if (byteBuffer == null) {
                 throw new NullPointerException("onNext must not be called with null byteBuffer");
             }
-            bodyPartPublisher.send(byteBuffer).whenComplete((r, t) -> {
+            publisherToUpstream.send(byteBuffer).whenComplete((r, t) -> {
                 if (t != null) {
                     handleError(t);
                     return;
@@ -290,7 +318,6 @@ public void onNext(ByteBuffer byteBuffer) {
 
         @Override
         public void onError(Throwable t) {
-            bodyPartPublisher.error(t);
             handleError(t);
         }
 
@@ -300,11 +327,11 @@ public void onComplete() {
         }
 
         private void handleError(Throwable t) {
+            publisherToUpstream.error(t);
             future.completeExceptionally(t);
         }
     }
 
-
     public static <ResponseT, ResultT> Builder<ResponseT, ResultT> builder() {
         return new Builder<>();
     }

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

@@ -18,10 +18,13 @@
 
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicInteger;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import org.reactivestreams.tck.SubscriberWhiteboxVerification;
 import org.reactivestreams.tck.TestEnvironment;
+import software.amazon.awssdk.core.ResponseBytes;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.utils.async.SimplePublisher;
 
 public class IndividualPartSubscriberTckTest extends SubscriberWhiteboxVerification<ByteBuffer> {
@@ -34,9 +37,15 @@ protected IndividualPartSubscriberTckTest() {
 
     @Override
     public Subscriber<ByteBuffer> createSubscriber(WhiteboxSubscriberProbe<ByteBuffer> probe) {
-        CompletableFuture<Integer> future = new CompletableFuture<>();
+        CompletableFuture<ByteBuffer> future = new CompletableFuture<>();
         SimplePublisher<ByteBuffer> publisher = new SimplePublisher<>();
-        return new SplittingTransformer.IndividualPartSubscriber<Integer>(future, 0, publisher) {
+        SplittingTransformer<Object, ResponseBytes<Object>> transformer =
+            SplittingTransformer.<Object, ResponseBytes<Object>>builder()
+                                .upstreamResponseTransformer(AsyncResponseTransformer.toBytes())
+                                .maximumBufferSizeInBytes(32L)
+                                .resultFuture(new CompletableFuture<>())
+                                .build();
+        return transformer.new IndividualPartSubscriber<ByteBuffer>(future, ByteBuffer.wrap(new byte[0])) {
             @Override
             public void onSubscribe(Subscription s) {
                 super.onSubscribe(s);

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

@@ -16,7 +16,6 @@
 package software.amazon.awssdk.services.s3.internal.multipart;
 
 import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
@@ -27,86 +26,136 @@
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.utils.Logger;
 
-// [WIP]
-// Still work in progress, currently only used to help manual testing, please ignore
+/**
+ * A subscriber implementation that will download all individual parts for a multipart get-object request. It receives the
+ * individual {@link AsyncResponseTransformer} which will be used to perform the individual part requests.
+ * This is a 'one-shot' class, it should <em>NOT</em> be reused for more than one multipart download
+ */
 @SdkInternalApi
 public class MultipartDownloaderSubscriber implements Subscriber<AsyncResponseTransformer<GetObjectResponse, GetObjectResponse>> {
     private static final Logger log = Logger.loggerFor(MultipartDownloaderSubscriber.class);
 
+    /**
+     * The s3 client used to make the individual part requests
+     */
     private final S3AsyncClient s3;
+
+    /**
+     * The GetObjectRequest that was provided when calling s3.getObject(...). It is copied for each individual request, and the
+     * copy has the partNumber field updated as more parts are downloaded.
+     */
     private final GetObjectRequest getObjectRequest;
 
-    private AtomicBoolean totalPartKnown = new AtomicBoolean(false);
-    private AtomicInteger totalParts = new AtomicInteger();
-    private AtomicInteger completed = new AtomicInteger(0);
-    private AtomicInteger currentPart = new AtomicInteger(0);
+    /**
+     * This value indicates the total number of parts of the object to get. If null, it means we don't know the total amount of
+     * parts, wither because we haven't received a response from s3 yet to set it, or the object to get is not multipart.
+     */
+    private volatile Integer totalParts;
+
+    /**
+     * The total number of completed parts. A part is considered complete once the completable future associated with its request
+     * completes successfully.
+     */
+    private final AtomicInteger completedParts = new AtomicInteger(0);
 
+    /**
+     * The subscription received from the publisher this subscriber subscribes to.
+     */
     private Subscription subscription;
 
-    private CompletableFuture<GetObjectResponse> responseFuture;
-    private GetObjectResponse returnResponse;
+    /**
+     * This future will be completed once this subscriber reaches a terminal state, failed or successfully, and will be
+     * completed accordingly.
+     */
+    private final CompletableFuture<Void> future = new CompletableFuture<>();
+
+    private final Object lock = new Object();
+
+    /**
+     * The etag of the object being downloaded.
+     */
+    private String eTag;
 
     public MultipartDownloaderSubscriber(S3AsyncClient s3, GetObjectRequest getObjectRequest) {
         this.s3 = s3;
         this.getObjectRequest = getObjectRequest;
-        this.responseFuture = new CompletableFuture<>();
     }
 
     @Override
     public void onSubscribe(Subscription s) {
-        if (this.subscription == null) {
-            this.subscription = s;
+        if (this.subscription != null) {
+            s.cancel();
+            return;
         }
-        s.request(1);
+        this.subscription = s;
+        this.subscription.request(1);
     }
 
     @Override
     public void onNext(AsyncResponseTransformer<GetObjectResponse, GetObjectResponse> asyncResponseTransformer) {
-        int part = currentPart.incrementAndGet();
-        if (totalPartKnown.get()) {
-            log.info(() -> String.format("total part: %s, current part: %s", totalParts.get(), part));
-        } else {
-            log.info(() -> String.format("part %s", part));
+        log.trace(() -> "onNext " + completedParts.get());
+        if (asyncResponseTransformer == null) {
+            throw new NullPointerException("onNext must not be called with null asyncResponseTransformer");
         }
-        if (totalPartKnown.get() && part > totalParts.get()) {
-            log.info(() -> "no more parts available, stopping");
-            subscription.cancel();
-            return;
+
+        int nextPartToGet = completedParts.get() + 1;
+
+        if (totalParts != null && nextPartToGet > totalParts) {
+            synchronized (lock) {
+                subscription.cancel();
+            }
         }
-        GetObjectRequest actualRequest = this.getObjectRequest.copy(req -> req.partNumber(part));
+
+        GetObjectRequest actualRequest = nextRequest(nextPartToGet);
         CompletableFuture<GetObjectResponse> future = s3.getObject(actualRequest, asyncResponseTransformer);
-        future.whenComplete((response, e) -> {
-            if (e != null) {
-                responseFuture.completeExceptionally(e);
+        future.whenComplete((response, error) -> {
+            if (error != null) {
+                onError(error);
                 return;
             }
-            completed.incrementAndGet();
-            returnResponse = response;
-            log.info(() -> String.format("received '%s'", response.contentRange()));
+            int totalComplete = completedParts.incrementAndGet();
+            log.trace(() -> String.format("completed part: %s", totalComplete));
+
+            if (eTag == null) {
+                this.eTag = response.eTag();
+                log.trace(() -> String.format("Multipart object ETag: %s", this.eTag));
+            }
+
             Integer partCount = response.partsCount();
-            if (totalPartKnown.compareAndSet(false, true)) {
-                totalParts.set(partCount);
-                totalPartKnown.set(true);
+            if (partCount != null && totalParts == null) {
+                log.trace(() -> String.format("total parts: %s", partCount));
+                totalParts = partCount;
             }
-            log.info(() -> String.format("total parts: %s", partCount));
-            if (totalParts.get() > 1) {
-                subscription.request(1);
+            synchronized (lock) {
+                if (totalParts != null && totalParts > 1 && totalComplete < totalParts) {
+                    subscription.request(1);
+                } else {
+                    subscription.cancel();
+                }
             }
         });
     }
 
     @Override
     public void onError(Throwable t) {
-        responseFuture.completeExceptionally(t);
+        future.completeExceptionally(t);
     }
 
     @Override
     public void onComplete() {
-        responseFuture.complete(returnResponse);
+        future.complete(null);
     }
 
-    public CompletableFuture<GetObjectResponse> future() {
-        return responseFuture;
+    public CompletableFuture<Void> future() {
+        return this.future;
     }
 
+    private GetObjectRequest nextRequest(int nextPartToGet) {
+        return getObjectRequest.copy(req -> {
+            req.partNumber(nextPartToGet);
+            if (eTag != null) {
+                req.ifMatch(eTag);
+            }
+        });
+    }
 }