Fix rare race condition in multipart upload with low maxInFlightParts

.changes/next-release/bugfix-AmazonS3-8f2c1a4.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon S3",
+    "contributor": "",
+    "description": "Fix a race condition in multipart upload with low maxInFlightParts where CompleteMultipartUpload could be initiated while the final part was still uploading."
+}

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

@@ -207,7 +207,7 @@ public void onNext(CloseableAsyncRequestBody asyncRequestBody) {
                                              subscription.request(1);
                                          }
                                      }
-                                     completeMultipartUploadIfFinished(inFlight);
+                                     completeMultipartUploadIfFinished();
                                  }
                              });
     }
@@ -263,12 +263,13 @@ public void onComplete() {
         log.debug(() -> "Received onComplete()");
         isDone = true;
         if (!isPaused) {
-            completeMultipartUploadIfFinished(asyncRequestBodyInFlight.get());
+            completeMultipartUploadIfFinished();
         }
     }
 
-    private void completeMultipartUploadIfFinished(int requestsInFlight) {
-        if (isDone && requestsInFlight == 0 && completedMultipartInitiated.compareAndSet(false, true)) {
+    private void completeMultipartUploadIfFinished() {
+        // All atomics including asyncRequestBodyInFlight MUST be re-read here rather than passed in by the caller.
+        if (isDone && asyncRequestBodyInFlight.get() == 0 && completedMultipartInitiated.compareAndSet(false, true)) {
             CompletedPart[] parts;
 
             if (existingParts.isEmpty()) {

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

@@ -232,7 +232,7 @@ private void sendUploadPartRequest(String uploadId,
                             subscription.request(1);
                         }
                     }
-                    completeMultipartUploadIfFinish(inFlight);
+                    completeMultipartUploadIfFinish();
                 }
             });
     }
@@ -266,12 +266,13 @@ public void onComplete() {
             multipartUploadHelper.uploadInOneChunk(putObjectRequest, entireRequestBody, returnFuture);
         } else {
             isDone = true;
-            completeMultipartUploadIfFinish(asyncRequestBodyInFlight.get());
+            completeMultipartUploadIfFinish();
         }
     }
 
-    private void completeMultipartUploadIfFinish(int requestsInFlight) {
-        if (isDone && requestsInFlight == 0 && completedMultipartInitiated.compareAndSet(false, true)) {
+    private void completeMultipartUploadIfFinish() {
+        // All atomics including asyncRequestBodyInFlight MUST be re-read here rather than passed in by the caller.
+        if (isDone && asyncRequestBodyInFlight.get() == 0 && completedMultipartInitiated.compareAndSet(false, true)) {
             CompletedPart[] parts = completedParts.stream()
                                                   .sorted(Comparator.comparingInt(CompletedPart::partNumber))
                                                   .toArray(CompletedPart[]::new);

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java

@@ -17,8 +17,10 @@
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -42,6 +44,8 @@
 import software.amazon.awssdk.core.async.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.internal.multipart.utils.ControlledSubscription;
+import software.amazon.awssdk.services.s3.internal.multipart.utils.ManagedUploadPart;
 import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadResponse;
 import software.amazon.awssdk.services.s3.model.CompletedPart;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
@@ -261,6 +265,63 @@ void maxInFlightPutObjectParts_shouldLimitConcurrentUploads() {
         verify(mockSubscription, times(1)).request(1);
     }
 
+    /**
+     * Regression test for early CompleteMultipartUpload with a missing part under low maxInFlightParts.
+     */
+    @Test
+    void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAllParts() {
+        int maxInFlight = 2;
+        int totalParts = 5;
+        long contentSize = totalParts * PART_SIZE;
+
+        MpuRequestContext context = MpuRequestContext.builder()
+                .request(Pair.of(putObjectRequest, asyncRequestBody))
+                .contentLength(contentSize)
+                .partSize(PART_SIZE)
+                .uploadId(UPLOAD_ID)
+                .numPartsCompleted(0L)
+                .expectedNumParts(totalParts)
+                .build();
+
+        ManagedUploadPart recorder = new ManagedUploadPart();
+        when(multipartUploadHelper.sendIndividualUploadPartRequest(eq(UPLOAD_ID), any(), any(), any(), any()))
+                .thenAnswer(recorder);
+
+        KnownContentLengthAsyncRequestBodySubscriber sub = createSubscriber(context, maxInFlight);
+        ControlledSubscription controlledSubscription = new ControlledSubscription(sub);
+        sub.onSubscribe(controlledSubscription);          // requests maxInFlight upfront
+
+        controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 1 starts
+        controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 2 starts
+        recorder.completePart(1);                                                          // frees a slot
+        controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 3 starts
+        recorder.completePart(2);
+        controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 4 starts
+
+        // Part 3 completes while the final chunk is not buffered yet: its request(1) delivers nothing.
+        recorder.completePart(3);
+
+        // The producer now queues the final body and the stream-complete signal. They will be delivered
+        // synchronously inside the next request(1), which happens in part 4's completion callback,
+        // between its decrementAndGet() (returning the stale 0) and completeMultipartUploadIfFinished.
+        controlledSubscription.enqueueBodyQuietly(createMockAsyncRequestBody(PART_SIZE));
+        controlledSubscription.enqueueStreamCompleteQuietly();
+        recorder.completePart(4);
+
+        // verify we haven't called CompleteMultipartUpload yet.
+        verify(multipartUploadHelper, never()).completeMultipartUpload(any(), any(), any(), any(), anyLong());
+        verify(multipartUploadHelper, never()).failRequestsElegantly(any(), any(), any(), any(), any());
+
+        // Part 5's upload finishes; only now should CompleteMultipartUpload be sent, with all 5 parts.
+        recorder.completePart(5);
+
+        ArgumentCaptor<CompletedPart[]> partsCaptor = ArgumentCaptor.forClass(CompletedPart[].class);
+        verify(multipartUploadHelper).completeMultipartUpload(eq(returnFuture), eq(UPLOAD_ID), partsCaptor.capture(),
+                                                              eq(putObjectRequest), eq(contentSize));
+        assertThat(partsCaptor.getValue()).hasSize(totalParts).doesNotContainNull();
+        verify(multipartUploadHelper, never()).failRequestsElegantly(any(), any(), any(), any(), any());
+    }
+
     private MpuRequestContext createDefaultMpuRequestContext() {
         return MpuRequestContext.builder()
                 .request(Pair.of(putObjectRequest, AsyncRequestBody.fromFile(testFile)))

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java

@@ -18,8 +18,10 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -32,7 +34,8 @@
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.core.async.CloseableAsyncRequestBody;
 import software.amazon.awssdk.core.exception.SdkClientException;
-import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.internal.multipart.utils.ControlledSubscription;
+import software.amazon.awssdk.services.s3.internal.multipart.utils.ManagedUploadPart;
 import software.amazon.awssdk.services.s3.model.CompletedPart;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
@@ -165,6 +168,57 @@ void onComplete_withNoParts_shouldUploadEmptyBody() {
         verify(multipartUploadHelper).uploadInOneChunk(eq(putObjectRequest), any(), eq(returnFuture));
     }
 
+    /**
+     * Regression test for the "Expected: N, Actual: N-1" part-count failure seen with low maxInFlightParts.
+     */
+    @Test
+    void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAllParts() {
+        int maxInFlight = 2;
+        int numParts = 5;
+        UnknownContentLengthAsyncRequestBodySubscriber subscriber = createSubscriber(maxInFlight);
+
+        stubSuccessfulCreateMultipartCall();
+        when(genericMultipartHelper.determinePartCount(anyLong(), anyLong()))
+            .thenAnswer(invocation -> (int) Math.ceil(invocation.getArgument(0, Long.class)
+                                                      / (double) invocation.getArgument(1, Long.class)));
+        ManagedUploadPart recorder = new ManagedUploadPart();
+        when(multipartUploadHelper.sendIndividualUploadPartRequest(any(), any(), any(), any(), any()))
+            .thenAnswer(recorder);
+
+        ControlledSubscription subscription = new ControlledSubscription(subscriber);
+        subscriber.onSubscribe(subscription);
+
+        subscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // held as firstRequestBody
+        subscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // triggers MPU; parts 1, 2 start
+        recorder.completePart(1);                                                  // frees a slot
+        subscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 3 starts
+        recorder.completePart(2);
+        subscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 4 starts
+
+        // Part 3 completes while the final chunk is not buffered yet: its request(1) delivers nothing.
+        recorder.completePart(3);
+
+        // The producer now queues the final body and the stream-complete signal. They will be delivered
+        // synchronously inside the next request(1), which happens in part 4's completion callback,
+        // between its decrementAndGet() (returning the stale 0) and completeMultipartUploadIfFinish.
+        subscription.enqueueBodyQuietly(createMockAsyncRequestBody(PART_SIZE));
+        subscription.enqueueStreamCompleteQuietly();
+        recorder.completePart(4);
+
+        // verify we haven't called CompleteMultipartUpload yet.
+        verify(multipartUploadHelper, never()).failRequestsElegantly(any(), any(), any(), any(), any());
+        verify(multipartUploadHelper, never()).completeMultipartUpload(any(), any(), any(), any(), anyLong());
+
+        // Part 5's upload finishes; only now should CompleteMultipartUpload be sent, with all 5 parts.
+        recorder.completePart(5);
+
+        ArgumentCaptor<CompletedPart[]> partsCaptor = ArgumentCaptor.forClass(CompletedPart[].class);
+        verify(multipartUploadHelper).completeMultipartUpload(eq(returnFuture), eq(UPLOAD_ID), partsCaptor.capture(),
+                                                              eq(putObjectRequest), eq(numParts * PART_SIZE));
+        assertThat(partsCaptor.getValue()).hasSize(numParts).doesNotContainNull();
+        verify(multipartUploadHelper, never()).failRequestsElegantly(any(), any(), any(), any(), any());
+    }
+
     private UnknownContentLengthAsyncRequestBodySubscriber createSubscriber(int maxInFlightParts) {
         return new UnknownContentLengthAsyncRequestBodySubscriber(
                 PART_SIZE, putObjectRequest, returnFuture,

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ControlledSubscription.java

@@ -0,0 +1,82 @@
+/*
+ * 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.services.s3.internal.multipart.utils;
+
+import java.util.ArrayDeque;
+import java.util.Deque;
+import org.reactivestreams.Subscriber;
+import org.reactivestreams.Subscription;
+import software.amazon.awssdk.core.async.CloseableAsyncRequestBody;
+
+/**
+ * A Subscription that mimics {@link software.amazon.awssdk.utils.async.SimplePublisher}: queued signals
+ * are delivered synchronously on the thread that calls request(), onNext delivery is gated on demand,
+ * and onComplete is delivered once the queue drains, without needing demand.
+ */
+public final class ControlledSubscription implements Subscription {
+    private final Subscriber<? super CloseableAsyncRequestBody> subscriber;
+    private final Deque<CloseableAsyncRequestBody> queuedBodies = new ArrayDeque<>();
+    private long demand;
+    private boolean streamComplete;
+    private boolean onCompleteDelivered;
+    private boolean draining;
+
+    public ControlledSubscription(Subscriber<? super CloseableAsyncRequestBody> subscriber) {
+        this.subscriber = subscriber;
+    }
+
+    @Override
+    public void request(long n) {
+        demand += n;
+        drain();
+    }
+
+    @Override
+    public void cancel() {
+    }
+
+    public void enqueueBodyAndDeliver(CloseableAsyncRequestBody body) {
+        queuedBodies.add(body);
+        drain();
+    }
+
+    public void enqueueBodyQuietly(CloseableAsyncRequestBody body) {
+        queuedBodies.add(body);
+    }
+
+    public void enqueueStreamCompleteQuietly() {
+        streamComplete = true;
+    }
+
+    private void drain() {
+        if (draining) {
+            return; // mirrors SimplePublisher's processingQueue flag: no re-entrant delivery
+        }
+        draining = true;
+        try {
+            while (demand > 0 && !queuedBodies.isEmpty()) {
+                demand--;
+                subscriber.onNext(queuedBodies.poll());
+            }
+            if (streamComplete && queuedBodies.isEmpty() && !onCompleteDelivered) {
+                onCompleteDelivered = true;
+                subscriber.onComplete();
+            }
+        } finally {
+            draining = false;
+        }
+    }
+}

services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ManagedUploadPart.java

@@ -0,0 +1,61 @@
+/*
+ * 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.services.s3.internal.multipart.utils;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.function.Consumer;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.services.s3.model.CompletedPart;
+import software.amazon.awssdk.services.s3.model.UploadPartRequest;
+import software.amazon.awssdk.utils.Pair;
+
+/**
+ * Records the consumer and pending future of each UploadPart request so a test can complete parts
+ * the same way MultipartUploadHelper does: consumer.accept(completedPart) followed by future completion.
+ *
+ * <p>Intended to be used as a Mockito {@link Answer} for
+ * {@code MultipartUploadHelper#sendIndividualUploadPartRequest}.
+ */
+public final class ManagedUploadPart implements Answer<CompletableFuture<CompletedPart>> {
+    private final Map<Integer, Consumer<CompletedPart>> consumers = new HashMap<>();
+    private final Map<Integer, CompletableFuture<CompletedPart>> futures = new HashMap<>();
+
+    @Override
+    @SuppressWarnings("unchecked")
+    public CompletableFuture<CompletedPart> answer(InvocationOnMock invocation) {
+        Consumer<CompletedPart> consumer = invocation.getArgument(1, Consumer.class);
+        Pair<UploadPartRequest, AsyncRequestBody> pair = invocation.getArgument(3, Pair.class);
+        int partNumber = pair.left().partNumber();
+        CompletableFuture<CompletedPart> future = new CompletableFuture<>();
+        consumers.put(partNumber, consumer);
+        futures.put(partNumber, future);
+        return future;
+    }
+
+    public void completePart(int partNumber) {
+        CompletableFuture<CompletedPart> future = futures.get(partNumber);
+        assertThat(future).withFailMessage("UploadPart request for part %d was never sent", partNumber).isNotNull();
+        CompletedPart part = CompletedPart.builder().partNumber(partNumber).eTag("etag-" + partNumber).build();
+        consumers.get(partNumber).accept(part);
+        future.complete(part);
+    }
+}