Fixed OOM issue that could surface when users try to download objects… (#5025)

* Fixed OOM issue that could surface when users try to download objects from an S3 bucket that has millions of small files.

* Fix sonatype issue

Author: zoewangg

.changes/next-release/bugfix-S3TransferManager-7a79dcf.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "S3 Transfer Manager",
+    "contributor": "",
+    "description": "Fixed OOM issue that could surface when users try to download objects from an S3 bucket that has millions of small files. See [#4987](https://github.com/aws/aws-sdk-java-v2/issues/4987)"
+}

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/AsyncBufferingSubscriber.java

@@ -15,7 +15,9 @@
 
 package software.amazon.awssdk.transfer.s3.internal;
 
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.Function;
 import org.reactivestreams.Subscriber;
@@ -40,13 +42,22 @@ public class AsyncBufferingSubscriber<T> implements Subscriber<T> {
     private volatile boolean upstreamDone;
     private Subscription subscription;
 
+    private final Set<CompletableFuture<?>> requestsInFlight;
+
     public AsyncBufferingSubscriber(Function<T, CompletableFuture<?>> consumer,
                                     CompletableFuture<Void> returnFuture,
                                     int maxConcurrentExecutions) {
         this.returnFuture = returnFuture;
         this.consumer = consumer;
         this.maxConcurrentExecutions = maxConcurrentExecutions;
         this.numRequestsInFlight = new AtomicInteger(0);
+        this.requestsInFlight = ConcurrentHashMap.newKeySet();
+
+        returnFuture.whenComplete((r, t) -> {
+            if (t != null) {
+                requestsInFlight.forEach(f -> f.cancel(true));
+            }
+        });
     }
 
     @Override
@@ -64,8 +75,11 @@ public void onSubscribe(Subscription subscription) {
     @Override
     public void onNext(T item) {
         numRequestsInFlight.incrementAndGet();
-        consumer.apply(item).whenComplete((r, t) -> {
+        CompletableFuture<?> currentRequest = consumer.apply(item);
+        requestsInFlight.add(currentRequest);
+        currentRequest.whenComplete((r, t) -> {
             checkForCompletion(numRequestsInFlight.decrementAndGet());
+            requestsInFlight.remove(currentRequest);
             synchronized (this) {
                 subscription.request(1);
             }

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/DownloadDirectoryHelper.java

@@ -106,15 +106,15 @@ private void doDownloadDirectory(CompletableFuture<CompletedDirectoryDownload> r
         Queue<FailedFileDownload> failedFileDownloads = new ConcurrentLinkedQueue<>();
 
         CompletableFuture<Void> allOfFutures = new CompletableFuture<>();
-
         AsyncBufferingSubscriber<S3Object> asyncBufferingSubscriber =
-            new AsyncBufferingSubscriber<>(downloadSingleFile(returnFuture, downloadDirectoryRequest, request,
+            new AsyncBufferingSubscriber<>(downloadSingleFile(downloadDirectoryRequest, request,
                                                               failedFileDownloads),
                                            allOfFutures,
                                            DEFAULT_DOWNLOAD_DIRECTORY_MAX_CONCURRENCY);
         listObjectsHelper.listS3ObjectsRecursively(request)
                          .filter(downloadDirectoryRequest.filter())
                          .subscribe(asyncBufferingSubscriber);
+        CompletableFutureUtils.forwardExceptionTo(returnFuture, allOfFutures);
 
         allOfFutures.whenComplete((r, t) -> {
             if (t != null) {
@@ -128,19 +128,14 @@ private void doDownloadDirectory(CompletableFuture<CompletedDirectoryDownload> r
     }
 
     private Function<S3Object, CompletableFuture<?>> downloadSingleFile(
-        CompletableFuture<CompletedDirectoryDownload> returnFuture,
         DownloadDirectoryRequest downloadDirectoryRequest,
         ListObjectsV2Request listRequest,
         Queue<FailedFileDownload> failedFileDownloads) {
 
-        return s3Object -> {
-            CompletableFuture<CompletedFileDownload> future = doDownloadSingleFile(downloadDirectoryRequest,
-                                                                                   failedFileDownloads,
-                                                                                   listRequest,
-                                                                                   s3Object);
-            CompletableFutureUtils.forwardExceptionTo(returnFuture, future);
-            return future;
-        };
+        return s3Object -> doDownloadSingleFile(downloadDirectoryRequest,
+                                            failedFileDownloads,
+                                            listRequest,
+                                            s3Object);
     }
 
     private Path determineDestinationPath(DownloadDirectoryRequest downloadDirectoryRequest,

services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/AsyncBufferingSubscriberTest.java

@@ -36,7 +36,7 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
-import org.reactivestreams.Subscription;
+import software.amazon.awssdk.utils.async.SimplePublisher;
 
 class AsyncBufferingSubscriberTest {
     private static final int MAX_CONCURRENT_EXECUTIONS = 5;
@@ -96,20 +96,16 @@ void differentNumberOfStrings_shouldCompleteSuccessfully(int numberOfStrings) th
     }
 
     @Test
-    void onErrorInvoked_shouldCompleteFutureExceptionally() {
-        subscriber.onSubscribe(new Subscription() {
-            @Override
-            public void request(long n) {
-
-            }
-
-            @Override
-            public void cancel() {
-
-            }
-        });
+    void onErrorInvoked_shouldCompleteFutureExceptionallyAndCancelRequestsFuture() {
         RuntimeException exception = new RuntimeException("test");
-        subscriber.onError(exception);
+        SimplePublisher<String> simplePublisher = new SimplePublisher<>();
+        simplePublisher.subscribe(subscriber);
+        simplePublisher.send("test");
+        simplePublisher.send("test");
+
+        simplePublisher.error(exception);
         assertThat(returnFuture).isCompletedExceptionally();
+        assertThat(futures.get(0)).isCancelled();
+        assertThat(futures.get(1)).isCancelled();
     }
 }

test/stability-tests/README.md

@@ -31,10 +31,9 @@ First add tests to TestRunner Class, then run the following command.
 
 ```
 mvn clean install -pl :stability-tests --am -P quick
-mvn clean install -pl :bom-inernal
-cd test/stability-tests
-mvn package -P test-jar
-java -jar target/stability-tests-uber.jar
+mvn clean install -pl :bom-internal
+mvn package -P test-jar -pl :stability-tests
+java -jar test/stability-tests/target/stability-tests-uber.jar
 ```
 
 ## Adding New Tests