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." +} 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..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 @@ -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()) { 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..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 @@ -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); 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..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 @@ -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 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))) 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..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 @@ -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 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, 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); + } +}