Add retry tests and onError cancel fix for presigned URL multipart downloads

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

@@ -16,7 +16,9 @@
 package software.amazon.awssdk.services.s3.internal.multipart;
 
 import java.util.Optional;
+import java.util.Queue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
@@ -71,6 +73,11 @@ public class PresignedUrlMultipartDownloaderSubscriber
     private final AtomicInteger completedParts;
     private final AtomicInteger requestsSent;
 
+    /**
+     * Store the GetObject futures so we can cancel them if onError() is invoked.
+     */
+    private final Queue<CompletableFuture<GetObjectResponse>> getObjectFutures = new ConcurrentLinkedQueue<>();
+
     private volatile Long totalContentLength;
     private volatile Integer totalParts;
     private volatile String eTag;
@@ -126,26 +133,27 @@ private void makeRangeRequest(int partIndex,
         log.debug(() -> "Sending range request for part " + partIndex + " with range=" + partRequest.range());
 
         requestsSent.incrementAndGet();
-        s3AsyncClient.presignedUrlExtension()
-                     .getObject(partRequest, asyncResponseTransformer)
-                     .whenComplete((response, error) -> {
-                         if (error != null) {
-                             if (partIndex == 0 && PresignedUrlDownloadHelper.isRangeNotSatisfiable(error)) {
-                                 log.debug(() -> "Received 416 on first range request, object is empty");
-                                 resultFuture.completeExceptionally(
-                                     new PresignedUrlDownloadHelper.EmptyObjectRangeNotSatisfiableException(error));
-                                 synchronized (lock) {
-                                     subscription.cancel();
-                                 }
-                             } else {
-                                 handleError(error);
-                             }
-                             return;
-                         }
-                         if (validatePart(response, partIndex, asyncResponseTransformer)) {
-                             requestMoreIfNeeded(completedParts.get());
-                         }
-                     });
+        CompletableFuture<GetObjectResponse> responseFuture = s3AsyncClient.presignedUrlExtension()
+                                                                           .getObject(partRequest, asyncResponseTransformer);
+        getObjectFutures.add(responseFuture);
+        responseFuture.whenComplete((response, error) -> {
+            if (error != null) {
+                if (partIndex == 0 && PresignedUrlDownloadHelper.isRangeNotSatisfiable(error)) {
+                    log.debug(() -> "Received 416 on first range request, object is empty");
+                    resultFuture.completeExceptionally(
+                        new PresignedUrlDownloadHelper.EmptyObjectRangeNotSatisfiableException(error));
+                    synchronized (lock) {
+                        subscription.cancel();
+                    }
+                } else {
+                    handleError(error);
+                }
+                return;
+            }
+            if (validatePart(response, partIndex, asyncResponseTransformer)) {
+                requestMoreIfNeeded(completedParts.get());
+            }
+        });
     }
 
     private boolean validatePart(GetObjectResponse response, int partIndex,
@@ -292,6 +300,10 @@ private void handleError(Throwable t) {
     @Override
     public void onError(Throwable t) {
         log.debug(() -> "Error in multipart download", t);
+        CompletableFuture<GetObjectResponse> partFuture;
+        while ((partFuture = getObjectFutures.poll()) != null) {
+            partFuture.cancel(true);
+        }
         future.completeExceptionally(t);
         if (resultFuture != null) {
             resultFuture.completeExceptionally(t);