From d0a0b5204b3400fd1f050695c11f456c229c63fc Mon Sep 17 00:00:00 2001 From: Alex Woods Date: Tue, 21 Jul 2026 09:43:40 -0700 Subject: [PATCH 1/3] 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. --- ...ntentLengthAsyncRequestBodySubscriber.java | 16 +- ...ntentLengthAsyncRequestBodySubscriber.java | 16 +- ...tLengthAsyncRequestBodySubscriberTest.java | 165 ++++++++++++++++++ ...tLengthAsyncRequestBodySubscriberTest.java | 162 +++++++++++++++++ 4 files changed, 351 insertions(+), 8 deletions(-) diff --git a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java index fddcf1cb843b..0b46b4e38cc0 100644 --- a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java +++ b/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,20 @@ 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() { + // asyncRequestBodyInFlight MUST be re-read here rather than passed in by the caller. In the upload + // completion callback, subscription.request(1) is invoked between decrementAndGet() and this method, + // which can synchronously deliver the final AsyncRequestBody (starting a new upload) followed by + // onComplete() (setting isDone). A stale snapshot of 0 taken before those deliveries would then + // initiate CompleteMultipartUpload while the final part is still in flight, completing the upload + // with a missing (null) part. Reading isDone (volatile) before the counter guarantees that, once + // isDone is observed as true, all onNext() increments are visible, so a fresh read of 0 means every + // upload that was started has finished and its CompletedPart is visible in completedParts. + if (isDone && asyncRequestBodyInFlight.get() == 0 && completedMultipartInitiated.compareAndSet(false, true)) { CompletedPart[] parts; if (existingParts.isEmpty()) { diff --git a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java index f264307f1984..9e466f36b5f9 100644 --- a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java +++ b/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,20 @@ 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() { + // asyncRequestBodyInFlight MUST be re-read here rather than passed in by the caller. In the upload + // completion callback, subscription.request(1) is invoked between decrementAndGet() and this method, + // which can synchronously deliver the final AsyncRequestBody (starting a new upload) followed by + // onComplete() (setting isDone). A stale snapshot of 0 taken before those deliveries would then + // initiate CompleteMultipartUpload while the final part is still in flight, completing the upload + // with a missing part. Reading isDone (volatile) before the counter guarantees that, once isDone is + // observed as true, all onNext() increments are visible, so a fresh read of 0 means every upload + // that was started has finished and its CompletedPart is visible in completedParts. + if (isDone && asyncRequestBodyInFlight.get() == 0 && completedMultipartInitiated.compareAndSet(false, true)) { CompletedPart[] parts = completedParts.stream() .sorted(Comparator.comparingInt(CompletedPart::partNumber)) .toArray(CompletedPart[]::new); diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java index 3c93a38b635b..762157306ebf 100644 --- a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java +++ b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java @@ -17,18 +17,24 @@ 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; import java.io.IOException; +import java.util.ArrayDeque; import java.util.Collection; +import java.util.Deque; +import java.util.HashMap; import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -46,6 +52,7 @@ import software.amazon.awssdk.services.s3.model.CompletedPart; import software.amazon.awssdk.services.s3.model.PutObjectRequest; import software.amazon.awssdk.services.s3.model.PutObjectResponse; +import software.amazon.awssdk.services.s3.model.UploadPartRequest; import software.amazon.awssdk.services.s3.multipart.S3ResumeToken; import software.amazon.awssdk.testutils.RandomTempFile; import software.amazon.awssdk.utils.Pair; @@ -261,6 +268,164 @@ void maxInFlightPutObjectParts_shouldLimitConcurrentUploads() { verify(mockSubscription, times(1)).request(1); } + /** + * Regression test for early CompleteMultipartUpload with a missing part under low maxInFlightParts. + * + *

In the upload completion callback, {@code subscription.request(1)} is invoked between + * {@code asyncRequestBodyInFlight.decrementAndGet()} and {@code completeMultipartUploadIfFinished}. + * The upstream SimplePublisher delivers queued signals synchronously on the requesting thread, so + * that request can deliver the final AsyncRequestBody (starting a new upload) followed by + * onComplete() (setting isDone) before the callback finishes. If completion is then decided using + * the stale decrement snapshot of 0, the part-count validation passes (partNumber was already + * incremented by the final onNext) and CompleteMultipartUpload is sent with a null entry for the + * still-in-flight final part. + * + *

This test scripts that exact delivery order with a Subscription that mimics SimplePublisher's + * synchronous, demand-gated delivery. + */ + @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(); + + UploadPartRecorder recorder = new UploadPartRecorder(); + when(multipartUploadHelper.sendIndividualUploadPartRequest(eq(UPLOAD_ID), any(), any(), any(), any())) + .thenAnswer(recorder); + + KnownContentLengthAsyncRequestBodySubscriber sub = createSubscriber(context, maxInFlight); + ScriptedSubscription scriptedSubscription = new ScriptedSubscription(sub); + sub.onSubscribe(scriptedSubscription); // requests maxInFlight upfront + + scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 1 starts + scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 2 starts + recorder.completePart(1); // frees a slot + scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 3 starts + recorder.completePart(2); + scriptedSubscription.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. + scriptedSubscription.enqueueBodyQuietly(createMockAsyncRequestBody(PART_SIZE)); + scriptedSubscription.enqueueStreamCompleteQuietly(); + recorder.completePart(4); + + // With the stale-snapshot bug, CompleteMultipartUpload was initiated here with parts[4] == null: + 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 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()); + } + + /** + * Records the consumer and pending future of each UploadPart request so the test can complete parts + * the same way MultipartUploadHelper does: consumer.accept(completedPart) followed by future completion. + */ + private static final class UploadPartRecorder implements org.mockito.stubbing.Answer> { + private final Map> consumers = new HashMap<>(); + private final Map> futures = new HashMap<>(); + + @Override + @SuppressWarnings("unchecked") + public CompletableFuture answer(org.mockito.invocation.InvocationOnMock invocation) { + Consumer consumer = invocation.getArgument(1, Consumer.class); + Pair pair = invocation.getArgument(3, Pair.class); + int partNumber = pair.left().partNumber(); + CompletableFuture future = new CompletableFuture<>(); + consumers.put(partNumber, consumer); + futures.put(partNumber, future); + return future; + } + + void completePart(int partNumber) { + CompletableFuture 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); + } + } + + /** + * 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. + */ + private static final class ScriptedSubscription implements Subscription { + private final KnownContentLengthAsyncRequestBodySubscriber subscriber; + private final Deque queuedBodies = new ArrayDeque<>(); + private long demand; + private boolean streamComplete; + private boolean onCompleteDelivered; + private boolean draining; + + ScriptedSubscription(KnownContentLengthAsyncRequestBodySubscriber subscriber) { + this.subscriber = subscriber; + } + + @Override + public void request(long n) { + demand += n; + drain(); + } + + @Override + public void cancel() { + } + + void enqueueBodyAndDeliver(CloseableAsyncRequestBody body) { + queuedBodies.add(body); + drain(); + } + + void enqueueBodyQuietly(CloseableAsyncRequestBody body) { + queuedBodies.add(body); + } + + 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; + } + } + } + private MpuRequestContext createDefaultMpuRequestContext() { return MpuRequestContext.builder() .request(Pair.of(putObjectRequest, AsyncRequestBody.fromFile(testFile))) diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java index d6cae74afa61..3bdc634c329f 100644 --- a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java +++ b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java @@ -18,24 +18,34 @@ 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; +import java.util.ArrayDeque; +import java.util.Deque; +import java.util.HashMap; +import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.function.Consumer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import org.reactivestreams.Subscription; +import software.amazon.awssdk.core.async.AsyncRequestBody; 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.model.CompletedPart; import software.amazon.awssdk.services.s3.model.PutObjectRequest; import software.amazon.awssdk.services.s3.model.PutObjectResponse; +import software.amazon.awssdk.services.s3.model.UploadPartRequest; +import software.amazon.awssdk.utils.Pair; public class UnknownContentLengthAsyncRequestBodySubscriberTest { @@ -165,6 +175,158 @@ 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. + * + *

In the upload completion callback, {@code subscription.request(1)} is invoked between + * {@code asyncRequestBodyInFlight.decrementAndGet()} and {@code completeMultipartUploadIfFinish}. + * The upstream SimplePublisher delivers queued signals synchronously on the requesting thread, so that + * request can deliver the final AsyncRequestBody (starting a new upload) followed by onComplete() + * (setting isDone) before the callback finishes. If completion is then decided using the stale + * decrement snapshot of 0, CompleteMultipartUpload is initiated while the final part is still in + * flight, and the upload fails with "The number of UploadParts requests is not equal to the expected + * number of parts" (or, without the part-count validation, would complete with a missing part). + * + *

This test scripts that exact delivery order with a Subscription that mimics SimplePublisher's + * synchronous, demand-gated delivery. + */ + @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))); + UploadPartRecorder recorder = new UploadPartRecorder(); + when(multipartUploadHelper.sendIndividualUploadPartRequest(any(), any(), any(), any(), any())) + .thenAnswer(recorder); + + ScriptedSubscription subscription = new ScriptedSubscription(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); + + // With the stale-snapshot bug, completion was initiated here with only 4 parts: + 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 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()); + } + + /** + * Records the consumer and pending future of each UploadPart request so the test can complete parts + * the same way MultipartUploadHelper does: consumer.accept(completedPart) followed by future completion. + */ + private static final class UploadPartRecorder implements org.mockito.stubbing.Answer> { + private final Map> consumers = new HashMap<>(); + private final Map> futures = new HashMap<>(); + + @Override + @SuppressWarnings("unchecked") + public CompletableFuture answer(org.mockito.invocation.InvocationOnMock invocation) { + Consumer consumer = invocation.getArgument(1, Consumer.class); + Pair pair = invocation.getArgument(3, Pair.class); + int partNumber = pair.left().partNumber(); + CompletableFuture future = new CompletableFuture<>(); + consumers.put(partNumber, consumer); + futures.put(partNumber, future); + return future; + } + + void completePart(int partNumber) { + CompletableFuture 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); + } + } + + /** + * 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. + */ + private static final class ScriptedSubscription implements Subscription { + private final UnknownContentLengthAsyncRequestBodySubscriber subscriber; + private final Deque queuedBodies = new ArrayDeque<>(); + private long demand; + private boolean streamComplete; + private boolean onCompleteDelivered; + private boolean draining; + + ScriptedSubscription(UnknownContentLengthAsyncRequestBodySubscriber subscriber) { + this.subscriber = subscriber; + } + + @Override + public void request(long n) { + demand += n; + drain(); + } + + @Override + public void cancel() { + } + + void enqueueBodyAndDeliver(CloseableAsyncRequestBody body) { + queuedBodies.add(body); + drain(); + } + + void enqueueBodyQuietly(CloseableAsyncRequestBody body) { + queuedBodies.add(body); + } + + 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; + } + } + } + private UnknownContentLengthAsyncRequestBodySubscriber createSubscriber(int maxInFlightParts) { return new UnknownContentLengthAsyncRequestBodySubscriber( PART_SIZE, putObjectRequest, returnFuture, From c1fadb28940d168563da21bc7b6d2e63c648ba76 Mon Sep 17 00:00:00 2001 From: Alex Woods Date: Tue, 21 Jul 2026 09:45:55 -0700 Subject: [PATCH 2/3] docs: Add changelog entry for multipart upload race fix --- .changes/next-release/bugfix-AmazonS3-8f2c1a4.json | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 .changes/next-release/bugfix-AmazonS3-8f2c1a4.json diff --git a/.changes/next-release/bugfix-AmazonS3-8f2c1a4.json b/.changes/next-release/bugfix-AmazonS3-8f2c1a4.json new file mode 100644 index 000000000000..2d1a0d49ee57 --- /dev/null +++ b/.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." +} From 5d8f4610297361fb205dbe0e5a57d030744262b9 Mon Sep 17 00:00:00 2001 From: Alex Woods Date: Tue, 21 Jul 2026 11:06:31 -0700 Subject: [PATCH 3/3] minor PR cleanups + refactor out common test util classes --- ...ntentLengthAsyncRequestBodySubscriber.java | 9 +- ...ntentLengthAsyncRequestBodySubscriber.java | 9 +- ...tLengthAsyncRequestBodySubscriberTest.java | 128 ++---------------- ...tLengthAsyncRequestBodySubscriberTest.java | 118 +--------------- .../utils/ControlledSubscription.java | 82 +++++++++++ .../multipart/utils/ManagedUploadPart.java | 61 +++++++++ 6 files changed, 162 insertions(+), 245 deletions(-) create mode 100644 services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ControlledSubscription.java create mode 100644 services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ManagedUploadPart.java diff --git a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java index 0b46b4e38cc0..eddc1cd9e078 100644 --- a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java +++ b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriber.java @@ -268,14 +268,7 @@ public void onComplete() { } private void completeMultipartUploadIfFinished() { - // asyncRequestBodyInFlight MUST be re-read here rather than passed in by the caller. In the upload - // completion callback, subscription.request(1) is invoked between decrementAndGet() and this method, - // which can synchronously deliver the final AsyncRequestBody (starting a new upload) followed by - // onComplete() (setting isDone). A stale snapshot of 0 taken before those deliveries would then - // initiate CompleteMultipartUpload while the final part is still in flight, completing the upload - // with a missing (null) part. Reading isDone (volatile) before the counter guarantees that, once - // isDone is observed as true, all onNext() increments are visible, so a fresh read of 0 means every - // upload that was started has finished and its CompletedPart is visible in completedParts. + // 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; diff --git a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java index 9e466f36b5f9..40ed3d7ecabb 100644 --- a/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java +++ b/services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriber.java @@ -271,14 +271,7 @@ public void onComplete() { } private void completeMultipartUploadIfFinish() { - // asyncRequestBodyInFlight MUST be re-read here rather than passed in by the caller. In the upload - // completion callback, subscription.request(1) is invoked between decrementAndGet() and this method, - // which can synchronously deliver the final AsyncRequestBody (starting a new upload) followed by - // onComplete() (setting isDone). A stale snapshot of 0 taken before those deliveries would then - // initiate CompleteMultipartUpload while the final part is still in flight, completing the upload - // with a missing part. Reading isDone (volatile) before the counter guarantees that, once isDone is - // observed as true, all onNext() increments are visible, so a fresh read of 0 means every upload - // that was started has finished and its CompletedPart is visible in completedParts. + // 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)) diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java index 762157306ebf..e458ce333dae 100644 --- a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java +++ b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/KnownContentLengthAsyncRequestBodySubscriberTest.java @@ -26,15 +26,11 @@ import static org.mockito.Mockito.when; import java.io.IOException; -import java.util.ArrayDeque; import java.util.Collection; -import java.util.Deque; -import java.util.HashMap; import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentLinkedQueue; -import java.util.function.Consumer; import java.util.function.Function; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -48,11 +44,12 @@ 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; import software.amazon.awssdk.services.s3.model.PutObjectResponse; -import software.amazon.awssdk.services.s3.model.UploadPartRequest; import software.amazon.awssdk.services.s3.multipart.S3ResumeToken; import software.amazon.awssdk.testutils.RandomTempFile; import software.amazon.awssdk.utils.Pair; @@ -270,18 +267,6 @@ void maxInFlightPutObjectParts_shouldLimitConcurrentUploads() { /** * Regression test for early CompleteMultipartUpload with a missing part under low maxInFlightParts. - * - *

In the upload completion callback, {@code subscription.request(1)} is invoked between - * {@code asyncRequestBodyInFlight.decrementAndGet()} and {@code completeMultipartUploadIfFinished}. - * The upstream SimplePublisher delivers queued signals synchronously on the requesting thread, so - * that request can deliver the final AsyncRequestBody (starting a new upload) followed by - * onComplete() (setting isDone) before the callback finishes. If completion is then decided using - * the stale decrement snapshot of 0, the part-count validation passes (partNumber was already - * incremented by the final onNext) and CompleteMultipartUpload is sent with a null entry for the - * still-in-flight final part. - * - *

This test scripts that exact delivery order with a Subscription that mimics SimplePublisher's - * synchronous, demand-gated delivery. */ @Test void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAllParts() { @@ -298,20 +283,20 @@ void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAl .expectedNumParts(totalParts) .build(); - UploadPartRecorder recorder = new UploadPartRecorder(); + ManagedUploadPart recorder = new ManagedUploadPart(); when(multipartUploadHelper.sendIndividualUploadPartRequest(eq(UPLOAD_ID), any(), any(), any(), any())) .thenAnswer(recorder); KnownContentLengthAsyncRequestBodySubscriber sub = createSubscriber(context, maxInFlight); - ScriptedSubscription scriptedSubscription = new ScriptedSubscription(sub); - sub.onSubscribe(scriptedSubscription); // requests maxInFlight upfront + ControlledSubscription controlledSubscription = new ControlledSubscription(sub); + sub.onSubscribe(controlledSubscription); // requests maxInFlight upfront - scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 1 starts - scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 2 starts + controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 1 starts + controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 2 starts recorder.completePart(1); // frees a slot - scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 3 starts + controlledSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 3 starts recorder.completePart(2); - scriptedSubscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // part 4 starts + 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); @@ -319,11 +304,11 @@ void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAl // 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. - scriptedSubscription.enqueueBodyQuietly(createMockAsyncRequestBody(PART_SIZE)); - scriptedSubscription.enqueueStreamCompleteQuietly(); + controlledSubscription.enqueueBodyQuietly(createMockAsyncRequestBody(PART_SIZE)); + controlledSubscription.enqueueStreamCompleteQuietly(); recorder.completePart(4); - // With the stale-snapshot bug, CompleteMultipartUpload was initiated here with parts[4] == null: + // 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()); @@ -337,95 +322,6 @@ void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAl verify(multipartUploadHelper, never()).failRequestsElegantly(any(), any(), any(), any(), any()); } - /** - * Records the consumer and pending future of each UploadPart request so the test can complete parts - * the same way MultipartUploadHelper does: consumer.accept(completedPart) followed by future completion. - */ - private static final class UploadPartRecorder implements org.mockito.stubbing.Answer> { - private final Map> consumers = new HashMap<>(); - private final Map> futures = new HashMap<>(); - - @Override - @SuppressWarnings("unchecked") - public CompletableFuture answer(org.mockito.invocation.InvocationOnMock invocation) { - Consumer consumer = invocation.getArgument(1, Consumer.class); - Pair pair = invocation.getArgument(3, Pair.class); - int partNumber = pair.left().partNumber(); - CompletableFuture future = new CompletableFuture<>(); - consumers.put(partNumber, consumer); - futures.put(partNumber, future); - return future; - } - - void completePart(int partNumber) { - CompletableFuture 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); - } - } - - /** - * 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. - */ - private static final class ScriptedSubscription implements Subscription { - private final KnownContentLengthAsyncRequestBodySubscriber subscriber; - private final Deque queuedBodies = new ArrayDeque<>(); - private long demand; - private boolean streamComplete; - private boolean onCompleteDelivered; - private boolean draining; - - ScriptedSubscription(KnownContentLengthAsyncRequestBodySubscriber subscriber) { - this.subscriber = subscriber; - } - - @Override - public void request(long n) { - demand += n; - drain(); - } - - @Override - public void cancel() { - } - - void enqueueBodyAndDeliver(CloseableAsyncRequestBody body) { - queuedBodies.add(body); - drain(); - } - - void enqueueBodyQuietly(CloseableAsyncRequestBody body) { - queuedBodies.add(body); - } - - 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; - } - } - } - private MpuRequestContext createDefaultMpuRequestContext() { return MpuRequestContext.builder() .request(Pair.of(putObjectRequest, AsyncRequestBody.fromFile(testFile))) diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java index 3bdc634c329f..22b5a7019a2b 100644 --- a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java +++ b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/UnknownContentLengthAsyncRequestBodySubscriberTest.java @@ -26,26 +26,19 @@ import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import java.util.ArrayDeque; -import java.util.Deque; -import java.util.HashMap; -import java.util.Map; import java.util.Optional; import java.util.concurrent.CompletableFuture; -import java.util.function.Consumer; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; import org.reactivestreams.Subscription; -import software.amazon.awssdk.core.async.AsyncRequestBody; 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; -import software.amazon.awssdk.services.s3.model.UploadPartRequest; -import software.amazon.awssdk.utils.Pair; public class UnknownContentLengthAsyncRequestBodySubscriberTest { @@ -177,18 +170,6 @@ void onComplete_withNoParts_shouldUploadEmptyBody() { /** * Regression test for the "Expected: N, Actual: N-1" part-count failure seen with low maxInFlightParts. - * - *

In the upload completion callback, {@code subscription.request(1)} is invoked between - * {@code asyncRequestBodyInFlight.decrementAndGet()} and {@code completeMultipartUploadIfFinish}. - * The upstream SimplePublisher delivers queued signals synchronously on the requesting thread, so that - * request can deliver the final AsyncRequestBody (starting a new upload) followed by onComplete() - * (setting isDone) before the callback finishes. If completion is then decided using the stale - * decrement snapshot of 0, CompleteMultipartUpload is initiated while the final part is still in - * flight, and the upload fails with "The number of UploadParts requests is not equal to the expected - * number of parts" (or, without the part-count validation, would complete with a missing part). - * - *

This test scripts that exact delivery order with a Subscription that mimics SimplePublisher's - * synchronous, demand-gated delivery. */ @Test void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAllParts() { @@ -200,11 +181,11 @@ void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAl when(genericMultipartHelper.determinePartCount(anyLong(), anyLong())) .thenAnswer(invocation -> (int) Math.ceil(invocation.getArgument(0, Long.class) / (double) invocation.getArgument(1, Long.class))); - UploadPartRecorder recorder = new UploadPartRecorder(); + ManagedUploadPart recorder = new ManagedUploadPart(); when(multipartUploadHelper.sendIndividualUploadPartRequest(any(), any(), any(), any(), any())) .thenAnswer(recorder); - ScriptedSubscription subscription = new ScriptedSubscription(subscriber); + ControlledSubscription subscription = new ControlledSubscription(subscriber); subscriber.onSubscribe(subscription); subscription.enqueueBodyAndDeliver(createMockAsyncRequestBody(PART_SIZE)); // held as firstRequestBody @@ -224,7 +205,7 @@ void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAl subscription.enqueueStreamCompleteQuietly(); recorder.completePart(4); - // With the stale-snapshot bug, completion was initiated here with only 4 parts: + // 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()); @@ -238,95 +219,6 @@ void lastBodyAndOnCompleteDeliveredInsideCompletionCallback_shouldCompleteWithAl verify(multipartUploadHelper, never()).failRequestsElegantly(any(), any(), any(), any(), any()); } - /** - * Records the consumer and pending future of each UploadPart request so the test can complete parts - * the same way MultipartUploadHelper does: consumer.accept(completedPart) followed by future completion. - */ - private static final class UploadPartRecorder implements org.mockito.stubbing.Answer> { - private final Map> consumers = new HashMap<>(); - private final Map> futures = new HashMap<>(); - - @Override - @SuppressWarnings("unchecked") - public CompletableFuture answer(org.mockito.invocation.InvocationOnMock invocation) { - Consumer consumer = invocation.getArgument(1, Consumer.class); - Pair pair = invocation.getArgument(3, Pair.class); - int partNumber = pair.left().partNumber(); - CompletableFuture future = new CompletableFuture<>(); - consumers.put(partNumber, consumer); - futures.put(partNumber, future); - return future; - } - - void completePart(int partNumber) { - CompletableFuture 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); - } - } - - /** - * 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. - */ - private static final class ScriptedSubscription implements Subscription { - private final UnknownContentLengthAsyncRequestBodySubscriber subscriber; - private final Deque queuedBodies = new ArrayDeque<>(); - private long demand; - private boolean streamComplete; - private boolean onCompleteDelivered; - private boolean draining; - - ScriptedSubscription(UnknownContentLengthAsyncRequestBodySubscriber subscriber) { - this.subscriber = subscriber; - } - - @Override - public void request(long n) { - demand += n; - drain(); - } - - @Override - public void cancel() { - } - - void enqueueBodyAndDeliver(CloseableAsyncRequestBody body) { - queuedBodies.add(body); - drain(); - } - - void enqueueBodyQuietly(CloseableAsyncRequestBody body) { - queuedBodies.add(body); - } - - 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; - } - } - } - private UnknownContentLengthAsyncRequestBodySubscriber createSubscriber(int maxInFlightParts) { return new UnknownContentLengthAsyncRequestBodySubscriber( PART_SIZE, putObjectRequest, returnFuture, diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ControlledSubscription.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ControlledSubscription.java new file mode 100644 index 000000000000..2e4ade734aea --- /dev/null +++ b/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 subscriber; + private final Deque queuedBodies = new ArrayDeque<>(); + private long demand; + private boolean streamComplete; + private boolean onCompleteDelivered; + private boolean draining; + + public ControlledSubscription(Subscriber 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; + } + } +} diff --git a/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ManagedUploadPart.java b/services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/multipart/utils/ManagedUploadPart.java new file mode 100644 index 000000000000..4a03b81c8d8e --- /dev/null +++ b/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. + * + *

Intended to be used as a Mockito {@link Answer} for + * {@code MultipartUploadHelper#sendIndividualUploadPartRequest}. + */ +public final class ManagedUploadPart implements Answer> { + private final Map> consumers = new HashMap<>(); + private final Map> futures = new HashMap<>(); + + @Override + @SuppressWarnings("unchecked") + public CompletableFuture answer(InvocationOnMock invocation) { + Consumer consumer = invocation.getArgument(1, Consumer.class); + Pair pair = invocation.getArgument(3, Pair.class); + int partNumber = pair.left().partNumber(); + CompletableFuture future = new CompletableFuture<>(); + consumers.put(partNumber, consumer); + futures.put(partNumber, future); + return future; + } + + public void completePart(int partNumber) { + CompletableFuture 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); + } +}