Address comments pt 1

davidh44 · davidh44 · commit 7f8803f46d85 · 2025-08-01T17:34:18.000-07:00
diff --git a/.changes/next-release/bugfix-AmazonS3-263fed5.json b/.changes/next-release/bugfix-AmazonS3-263fed5.json
@@ -2,5 +2,5 @@
     "type": "bugfix",
     "category": "Amazon S3",
     "contributor": "",
-    "description": "Fix MultipartS3AsyncClient GetObject retryable errors"
+    "description": "Fix bug in MultipartS3AsyncClient GET where retryable errors may not retried, and if retried, successful responses are incorrectly processed with the initial error."
 }
diff --git a/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/SplittingTransformer.java b/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/SplittingTransformer.java
@@ -18,6 +18,7 @@
 import java.nio.ByteBuffer;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
@@ -57,7 +58,7 @@ public class SplittingTransformer<ResponseT, ResultT> implements SdkPublisher<As
     /**
      * Set to true once {@code .onStream()} is called on the upstreamResponseTransformer
      */
-    private final AtomicBoolean onStreamCalled = new AtomicBoolean(false);
+    private boolean onStreamCalled;
 
     /**
      * Set to true once {@code .cancel()} is called in the subscription of the downstream subscriber, or if the
@@ -108,16 +109,9 @@ public class SplittingTransformer<ResponseT, ResultT> implements SdkPublisher<As
     private volatile CompletableFuture<ResultT> upstreamFuture;
 
     /**
-     * Tracks whether an {@code IndividualTransformer} has been created for the first part yet. Errors will only be retried for
-     * the first part.
+     * Tracks the part number. Errors will only be retried for the first part.
      */
-    private final AtomicBoolean isFirstIndividualTransformer = new AtomicBoolean(true);
-
-    /**
-     * Tracks whether an {@code IndividualPartSubscriber} has been created for the first part yet. Errors will only be retried for
-     * the first part.
-     */
-    private final AtomicBoolean isFirstIndividualSubscriber = new AtomicBoolean(true);
+    private final AtomicInteger partNumber = new AtomicInteger(0);
 
     private SplittingTransformer(AsyncResponseTransformer<ResponseT, ResultT> upstreamResponseTransformer,
                                  Long maximumBufferSizeInBytes,
@@ -206,7 +200,7 @@ private boolean doEmit() {
             }
             if (outstandingDemand.get() > 0) {
                 demand = outstandingDemand.decrementAndGet();
-                downstreamSubscriber.onNext(new IndividualTransformer(isFirstIndividualTransformer.compareAndSet(true, false)));
+                downstreamSubscriber.onNext(new IndividualTransformer(partNumber.incrementAndGet()));
             }
         }
         return false;
@@ -224,7 +218,7 @@ private void handleSubscriptionCancel() {
                 log.trace(() -> "downstreamSubscriber already null, skipping downstreamSubscriber.onComplete()");
                 return;
             }
-            if (!onStreamCalled.get()) {
+            if (!onStreamCalled) {
                 // we never subscribe publisherToUpstream to the upstream, it would not complete
                 downstreamSubscriber = null;
                 return;
@@ -268,19 +262,19 @@ private void handleFutureCancel(Throwable e) {
      * body publisher.
      */
     private class IndividualTransformer implements AsyncResponseTransformer<ResponseT, ResponseT> {
-        private final boolean isFirstPart;
+        private final int partNumber;
         private ResponseT response;
         private CompletableFuture<ResponseT> individualFuture;
 
-        IndividualTransformer(boolean isFirstPart) {
-            this.isFirstPart = isFirstPart;
+        IndividualTransformer(int partNumber) {
+            this.partNumber = partNumber;
         }
 
         @Override
         public CompletableFuture<ResponseT> prepare() {
             this.individualFuture = new CompletableFuture<>();
 
-            if (isFirstPart) {
+            if (partNumber == 1) {
                 if (isCancelled.get()) {
                     return individualFuture;
                 }
@@ -299,7 +293,7 @@ public CompletableFuture<ResponseT> prepare() {
 
         @Override
         public void onResponse(ResponseT response) {
-            if (isFirstPart) {
+            if (partNumber == 1) {
                 log.trace(() -> "calling onResponse on the upstream transformer");
                 upstreamResponseTransformer.onResponse(response);
             }
@@ -312,8 +306,8 @@ public void onStream(SdkPublisher<ByteBuffer> publisher) {
                 return;
             }
             synchronized (cancelLock) {
-                if (isFirstPart) {
-                    onStreamCalled.set(true);
+                if (partNumber == 1) {
+                    onStreamCalled = true;
                     log.trace(() -> "calling onStream on the upstream transformer");
                     upstreamResponseTransformer.onStream(upstreamSubscriber -> publisherToUpstream.subscribe(
                         DelegatingBufferingSubscriber.builder()
@@ -324,22 +318,25 @@ public void onStream(SdkPublisher<ByteBuffer> publisher) {
                 }
             }
 
-            if (!resultFuture.isDone()) {
-                CompletableFutureUtils.forwardResultTo(upstreamFuture, resultFuture);
-            }
-
-            publisher.subscribe(new IndividualPartSubscriber<>(this.individualFuture, response,
-                                                               isFirstIndividualSubscriber.compareAndSet(true, false)));
+            CompletableFutureUtils.forwardResultTo(upstreamFuture, resultFuture);
+            publisher.subscribe(new IndividualPartSubscriber<>(this.individualFuture, response, partNumber));
         }
 
         @Override
         public void exceptionOccurred(Throwable error) {
             log.trace(() -> "calling exceptionOccurred on the upstream transformer");
-            upstreamResponseTransformer.exceptionOccurred(error);
 
-            if (!isFirstPart || onStreamCalled.get()) {
-                publisherToUpstream.error(error);
+            if (partNumber == 1) {
+                upstreamResponseTransformer.exceptionOccurred(error);
             }
+
+            // TODO - add comments explaining
+            synchronized (cancelLock) {
+                if (partNumber > 1 || onStreamCalled) {
+                    publisherToUpstream.error(error);
+                }
+            }
+
         }
     }
 
@@ -350,13 +347,13 @@ class IndividualPartSubscriber<T> implements Subscriber<ByteBuffer> {
 
         private final CompletableFuture<T> future;
         private final T response;
-        private final boolean isFirstPart;
+        private final int partNumber;
         private Subscription subscription;
 
-        IndividualPartSubscriber(CompletableFuture<T> future, T response, boolean isFirstPart) {
+        IndividualPartSubscriber(CompletableFuture<T> future, T response, int partNumber) {
             this.future = future;
             this.response = response;
-            this.isFirstPart = isFirstPart;
+            this.partNumber = partNumber;
         }
 
         @Override
@@ -396,7 +393,7 @@ public void onComplete() {
         }
 
         private void handleError(Throwable t) {
-            if (!isFirstPart) {
+            if (partNumber > 1) {
                 publisherToUpstream.error(t);
             }
             future.completeExceptionally(t);
diff --git a/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/IndividualPartSubscriberTckTest.java b/core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/IndividualPartSubscriberTckTest.java
@@ -44,7 +44,7 @@ public Subscriber<ByteBuffer> createSubscriber(WhiteboxSubscriberProbe<ByteBuffe
                                 .maximumBufferSizeInBytes(32L)
                                 .resultFuture(new CompletableFuture<>())
                                 .build();
-        return transformer.new IndividualPartSubscriber<ByteBuffer>(future, ByteBuffer.wrap(new byte[0]), true) {
+        return transformer.new IndividualPartSubscriber<ByteBuffer>(future, ByteBuffer.wrap(new byte[0]), 1) {
             @Override
             public void onSubscribe(Subscription s) {
                 super.onSubscribe(s);
diff --git a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartDownloaderSubscriber.java b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/MultipartDownloaderSubscriber.java
@@ -15,9 +15,9 @@
 
 package software.amazon.awssdk.services.s3.internal.multipart;
 
-import java.util.ArrayList;
-import java.util.List;
+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;
@@ -81,7 +81,7 @@ public class MultipartDownloaderSubscriber implements Subscriber<AsyncResponseTr
      */
     private final Object lock = new Object();
 
-    private final List<CompletableFuture<GetObjectResponse>> getObjectFutures = new ArrayList<>();
+    private final Queue<CompletableFuture<GetObjectResponse>> getObjectFutures = new ConcurrentLinkedQueue<>();
 
     public MultipartDownloaderSubscriber(S3AsyncClient s3, GetObjectRequest getObjectRequest) {
         this(s3, getObjectRequest, 0);
@@ -125,6 +125,7 @@ public void onNext(AsyncResponseTransformer<GetObjectResponse, GetObjectResponse
         CompletableFuture<GetObjectResponse> getObjectFuture = s3.getObject(actualRequest, asyncResponseTransformer);
         getObjectFutures.add(getObjectFuture);
         getObjectFuture.whenComplete((response, error) -> {
+            getObjectFutures.remove(getObjectFuture);
             if (error != null) {
                 log.debug(() -> "Error encountered during GetObjectRequest with partNumber=" + nextPartToGet);
                 onError(error);
@@ -171,7 +172,10 @@ private void requestMoreIfNeeded(GetObjectResponse response) {
 
     @Override
     public void onError(Throwable t) {
-        getObjectFutures.forEach(future -> future.cancel(true));
+        CompletableFuture<GetObjectResponse> partFuture;
+        while ((partFuture = getObjectFutures.poll()) != null) {
+            partFuture.cancel(true);
+        }
         future.completeExceptionally(t);
     }
 
diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/S3MultipartClientGetObjectWiremockTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/S3MultipartClientGetObjectWiremockTest.java
@@ -92,6 +92,8 @@ public void setup(WireMockRuntimeInfo wm) {
                                        .build();
     }
 
+    // TODO - 1) update test names 2) add test for I/O error
+
     @Test
     public void stub_200_only() {
         List<CompletableFuture<ResponseBytes<GetObjectResponse>>> futures = new ArrayList<>();
@@ -305,7 +307,7 @@ public void multipleParts_503OnFirstPart_then_200s() {
     }
 
     private CompletableFuture<ResponseBytes<GetObjectResponse>> mock200Response(S3AsyncClient s3Client, int runNumber) {
-        String runId = runNumber + " sucess";
+        String runId = runNumber + " success";
 
         stubFor(any(anyUrl())
                     .withHeader("RunNum", matching(runId))

Original file line number	Diff line number	Diff line change
`@@ -2,5 +2,5 @@`
`2`	`2`	`"type": "bugfix",`
`3`	`3`	`"category": "Amazon S3",`
`4`	`4`	`"contributor": "",`
`5`		`- "description": "Fix MultipartS3AsyncClient GetObject retryable errors"`
	`5`	`+ "description": "Fix bug in MultipartS3AsyncClient GET where retryable errors may not retried, and if retried, successful responses are incorrectly processed with the initial error."`
`6`	`6`	`}`
Original file line number	Diff line number	Diff line change
`@@ -92,6 +92,8 @@ public void setup(WireMockRuntimeInfo wm) {`
`92`	`92`	`.build();`
`93`	`93`	`}`
`94`	`94`
	`95`	`+ // TODO - 1) update test names 2) add test for I/O error`
	`96`	`+`
`95`	`97`	`@Test`
`96`	`98`	`public void stub_200_only() {`
`97`	`99`	`List<CompletableFuture<ResponseBytes<GetObjectResponse>>> futures = new ArrayList<>();`
`@@ -305,7 +307,7 @@ public void multipleParts_503OnFirstPart_then_200s() {`
`305`	`307`	`}`
`306`	`308`
`307`	`309`	`private CompletableFuture<ResponseBytes<GetObjectResponse>> mock200Response(S3AsyncClient s3Client, int runNumber) {`
`308`		`- String runId = runNumber + " sucess";`
	`310`	`+ String runId = runNumber + " success";`
`309`	`311`
`310`	`312`	`stubFor(any(anyUrl())`
`311`	`313`	`.withHeader("RunNum", matching(runId))`