Parallel split for multipart GetObject File Download (#6425)

Parallel split for multipart GetObject File Download

Author: L-Applin

.changes/next-release/feature-S3-12a0b55.json

@@ -0,0 +1,6 @@
+{
+    "type": "feature",
+    "category": "S3",
+    "contributor": "",
+    "description": "Add support for parallel download for individual part-get for multipart GetObject in s3 async client and Transfer Manager"
+}

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

@@ -381,6 +381,15 @@ interface SplitResult<ResponseT, ResultT>
          */
         CompletableFuture<ResultT> resultFuture();
 
+        /**
+         * Indicates if the split async response transformer supports sending individual transformer non-serially and
+         * receiving back data from the many {@link AsyncResponseTransformer#onStream(SdkPublisher) publishers} non-serially.
+         * @return true if non-serial data is supported, false otherwise
+         */
+        default Boolean parallelSplitSupported() {
+            return false;
+        }
+
         static <ResponseT, ResultT> Builder<ResponseT, ResultT> builder() {
             return DefaultAsyncResponseTransformerSplitResult.builder();
         }
@@ -413,6 +422,20 @@ interface Builder<ResponseT, ResultT>
              * @return an instance of this Builder
              */
             Builder<ResponseT, ResultT> resultFuture(CompletableFuture<ResultT> future);
+
+            /**
+             * If the AsyncResponseTransformers returned by the {@link SplitResult#publisher()} support concurrent
+             * parallel streaming of multiple content body concurrently.
+             * @return
+             */
+            Boolean parallelSplitSupported();
+
+            /**
+             * Sets whether the AsyncResponseTransformers returned by the {@link SplitResult#publisher()} support concurrent
+             * parallel streaming of multiple content body concurrently
+             * @return
+             */
+            Builder<ResponseT, ResultT> parallelSplitSupported(Boolean parallelSplitSupported);
         }
     }
 

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

@@ -20,6 +20,7 @@
 import org.reactivestreams.Subscriber;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
+import software.amazon.awssdk.core.SplittingTransformerConfiguration;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.utils.Logger;
@@ -108,6 +109,11 @@ public String name() {
             return delegate.name();
         }
 
+        @Override
+        public SplitResult<ResponseT, ResultT> split(SplittingTransformerConfiguration splitConfig) {
+            return delegate.split(splitConfig);
+        }
+
         static void invoke(Runnable runnable, String callbackName) {
             try {
                 runnable.run();

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

@@ -27,12 +27,14 @@ public final class DefaultAsyncResponseTransformerSplitResult<ResponseT, ResultT
 
     private final SdkPublisher<AsyncResponseTransformer<ResponseT, ResponseT>> publisher;
     private final CompletableFuture<ResultT> future;
+    private final Boolean parallelSplitSupported;
 
     private DefaultAsyncResponseTransformerSplitResult(Builder<ResponseT, ResultT> builder) {
         this.publisher = Validate.paramNotNull(
             builder.publisher(), "asyncResponseTransformerPublisher");
         this.future = Validate.paramNotNull(
             builder.resultFuture(), "future");
+        this.parallelSplitSupported = Validate.getOrDefault(builder.parallelSplitSupported(), () -> false);
     }
 
     /**
@@ -52,6 +54,11 @@ public CompletableFuture<ResultT> resultFuture() {
         return this.future;
     }
 
+    @Override
+    public Boolean parallelSplitSupported() {
+        return this.parallelSplitSupported;
+    }
+
     @Override
     public AsyncResponseTransformer.SplitResult.Builder<ResponseT, ResultT> toBuilder() {
         return new DefaultBuilder<>(this);
@@ -65,13 +72,15 @@ public static class DefaultBuilder<ResponseT, ResultT>
         implements AsyncResponseTransformer.SplitResult.Builder<ResponseT, ResultT> {
         private SdkPublisher<AsyncResponseTransformer<ResponseT, ResponseT>> publisher;
         private CompletableFuture<ResultT> future;
+        private Boolean parallelSplitSupported;
 
         DefaultBuilder() {
         }
 
         DefaultBuilder(DefaultAsyncResponseTransformerSplitResult<ResponseT, ResultT> split) {
             this.publisher = split.publisher;
             this.future = split.future;
+            this.parallelSplitSupported = split.parallelSplitSupported;
         }
 
         @Override
@@ -92,14 +101,28 @@ public CompletableFuture<ResultT> resultFuture() {
         }
 
         @Override
-        public AsyncResponseTransformer.SplitResult.Builder<ResponseT, ResultT> resultFuture(CompletableFuture<ResultT> future) {
+        public AsyncResponseTransformer.SplitResult.Builder<ResponseT, ResultT> resultFuture(
+            CompletableFuture<ResultT> future) {
             this.future = future;
             return this;
         }
 
+        @Override
+        public Boolean parallelSplitSupported() {
+            return parallelSplitSupported;
+        }
+
+        @Override
+        public AsyncResponseTransformer.SplitResult.Builder<ResponseT, ResultT> parallelSplitSupported(
+            Boolean parallelSplitSupported) {
+            this.parallelSplitSupported = parallelSplitSupported;
+            return this;
+        }
+
         @Override
         public AsyncResponseTransformer.SplitResult<ResponseT, ResultT> build() {
             return new DefaultAsyncResponseTransformerSplitResult<>(this);
         }
+
     }
 }

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

@@ -0,0 +1,144 @@
+/*
+ * 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.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Supplier;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.ThreadSafe;
+import software.amazon.awssdk.utils.Logger;
+
+/**
+ * Subscription which can emit {@link Subscriber#onNext(T)} signals to a subscriber, based on the demand received with the
+ * {@link Subscription#request(long)}. It tracks the outstandingDemand that has not yet been fulfilled and used a Supplier
+ * passed to it to create the object it needs to emit.
+ * @param <T> the type of object to emit to the subscriber.
+ */
+@SdkInternalApi
+@ThreadSafe
+public final class EmittingSubscription<T> implements Subscription {
+    private static final Logger log = Logger.loggerFor(EmittingSubscription.class);
+
+    private Subscriber<? super T> downstreamSubscriber;
+    private final AtomicBoolean emitting;
+    private final AtomicLong outstandingDemand;
+    private final Runnable onCancel;
+    private final AtomicBoolean isCancelled;
+    private final Supplier<T> supplier;
+
+    private EmittingSubscription(Builder<T> builder) {
+        this.downstreamSubscriber = builder.downstreamSubscriber;
+        this.onCancel = builder.onCancel;
+        this.supplier = builder.supplier;
+        this.isCancelled = new AtomicBoolean();
+        this.outstandingDemand = new AtomicLong(0);
+        this.emitting = new AtomicBoolean();
+    }
+
+    public static <T> Builder<T> builder() {
+        return new Builder<>();
+    }
+
+    @Override
+    public void request(long n) {
+        if (n <= 0) {
+            downstreamSubscriber.onError(new IllegalArgumentException("Amount requested must be positive"));
+            return;
+        }
+        long newDemand = outstandingDemand.updateAndGet(current -> {
+            if (Long.MAX_VALUE - current < n) {
+                return Long.MAX_VALUE;
+            }
+            return current + n;
+        });
+        log.trace(() -> String.format("new outstanding demand: %s", newDemand));
+        emit();
+    }
+
+    @Override
+    public void cancel() {
+        isCancelled.set(true);
+        downstreamSubscriber = null;
+        onCancel.run();
+    }
+
+    private void emit() {
+        do {
+            if (!emitting.compareAndSet(false, true)) {
+                return;
+            }
+            try {
+                if (doEmit()) {
+                    return;
+                }
+            } finally {
+                emitting.compareAndSet(true, false);
+            }
+        } while (outstandingDemand.get() > 0);
+    }
+
+    private boolean doEmit() {
+        long demand = outstandingDemand.get();
+
+        while (demand > 0) {
+            if (isCancelled.get()) {
+                return true;
+            }
+            if (outstandingDemand.get() > 0) {
+                demand = outstandingDemand.decrementAndGet();
+                T value;
+                try {
+                    value = supplier.get();
+                } catch (Exception e) {
+                    downstreamSubscriber.onError(e);
+                    return true;
+                }
+                downstreamSubscriber.onNext(value);
+            }
+        }
+        return false;
+    }
+
+    public static class Builder<T> {
+        private Subscriber<? super T> downstreamSubscriber;
+        private Runnable onCancel;
+        private Supplier<T> supplier;
+
+        public Builder<T> downstreamSubscriber(Subscriber<? super T> subscriber) {
+            this.downstreamSubscriber = subscriber;
+            return this;
+        }
+
+        public Builder<T> onCancel(Runnable onCancel) {
+            this.onCancel = onCancel;
+            return this;
+        }
+
+        public Builder<T> supplier(Supplier<T> supplier) {
+            this.supplier = supplier;
+            return this;
+        }
+
+        public EmittingSubscription<T> build() {
+            return new EmittingSubscription<>(this);
+        }
+    }
+
+
+}

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

@@ -41,6 +41,7 @@
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.FileTransformerConfiguration;
 import software.amazon.awssdk.core.FileTransformerConfiguration.FailureBehavior;
+import software.amazon.awssdk.core.SplittingTransformerConfiguration;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.SdkPublisher;
 import software.amazon.awssdk.core.exception.SdkClientException;
@@ -76,6 +77,18 @@ private FileAsyncResponseTransformer(Path path, FileTransformerConfiguration fil
         this.position = position;
     }
 
+    FileTransformerConfiguration config() {
+        return configuration.toBuilder().build();
+    }
+
+    Path path() {
+        return path;
+    }
+
+    long position() {
+        return position;
+    }
+
     private static long determineFilePositionToWrite(Path path, FileTransformerConfiguration fileConfiguration) {
         if (fileConfiguration.fileWriteOption() == CREATE_OR_APPEND_TO_EXISTING) {
             try {
@@ -89,7 +102,7 @@ private static long determineFilePositionToWrite(Path path, FileTransformerConfi
         if (fileConfiguration.fileWriteOption() == WRITE_TO_POSITION) {
             return Validate.getOrDefault(fileConfiguration.position(), () -> 0L);
         }
-        return  0L;
+        return 0L;
     }
 
     private AsynchronousFileChannel createChannel(Path path) throws IOException {
@@ -183,6 +196,7 @@ static class FileSubscriber implements Subscriber<ByteBuffer> {
         private final Path path;
         private final CompletableFuture<Void> future;
         private final Consumer<Throwable> onErrorMethod;
+        private final Object closeLock = new Object();
 
         private volatile boolean writeInProgress = false;
         private volatile boolean closeOnLastWrite = false;
@@ -228,7 +242,7 @@ public void completed(Integer result, ByteBuffer attachment) {
                     if (byteBuffer.hasRemaining()) {
                         performWrite(byteBuffer);
                     } else {
-                        synchronized (FileSubscriber.this) {
+                        synchronized (closeLock) {
                             writeInProgress = false;
                             if (closeOnLastWrite) {
                                 close();
@@ -256,7 +270,7 @@ public void onError(Throwable t) {
         public void onComplete() {
             log.trace(() -> "onComplete");
             // if write in progress, tell write to close on finish.
-            synchronized (this) {
+            synchronized (closeLock) {
                 if (writeInProgress) {
                     log.trace(() -> "writeInProgress = true, not closing");
                     closeOnLastWrite = true;
@@ -284,4 +298,18 @@ public String toString() {
             return getClass() + ":" + path.toString();
         }
     }
-}
+
+
+    @Override
+    public SplitResult<ResponseT, ResponseT> split(SplittingTransformerConfiguration splitConfig) {
+        if (configuration.fileWriteOption() == CREATE_OR_APPEND_TO_EXISTING) {
+            return AsyncResponseTransformer.super.split(splitConfig);
+        }
+        CompletableFuture<ResponseT> future = new CompletableFuture<>();
+        return (SplitResult<ResponseT, ResponseT>) SplitResult.<ResponseT, ResponseT>builder()
+                                                              .publisher(new FileAsyncResponseTransformerPublisher(this))
+                                                              .resultFuture(future)
+                                                              .parallelSplitSupported(true)
+                                                              .build();
+    }
+}