Fix rare race condition in multipart upload with low maxInFlightParts (#7172)

* fix(s3): Prevent premature CompleteMultipartUpload with low maxInFlightParts The multipart upload subscribers decided whether to initiate CompleteMultipartUpload using a stale snapshot of the in-flight part counter taken by decrementAndGet() in the upload completion callback. Between that snapshot and the completion check, subscription.request(1) can synchronously deliver the final AsyncRequestBody (starting a new upload) followed by onComplete() (setting isDone), because SimplePublisher delivers queued signals on the requesting thread. The callback then saw isDone == true with its stale count of 0 and initiated CompleteMultipartUpload while the final part was still in flight. For unknown content length this failed uploads with 'The number of UploadParts requests is not equal to the expected number of parts. Expected: N, Actual: N-1'. For known content length the part-count validation passed (partNumber was already incremented by the final onNext) and CompleteMultipartUpload was sent with a null part entry. The window only opens when delivery of the final body is gated on completion callbacks' request(1) calls, which is why it was observed with maxInFlightParts=2 but not with the default of 50, and only intermittently under real network timing. Fix: re-read asyncRequestBodyInFlight inside the completion check instead of trusting the caller's snapshot. Once the volatile isDone is observed true, all onNext increments are visible, so a fresh read of 0 guarantees every started upload has finished. Both regression tests script the exact delivery order with a SimplePublisher-faithful Subscription and fail without the fix.

* docs: Add changelog entry for multipart upload race fix

* minor PR cleanups + refactor out common test util classes

Author: alextwoods

.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);
+    }
+}