Fix bug in FileAsyncRequestBody where inflight parts were negative

Author: alextwoods

.changes/next-release/bugfix-AmazonS3-4025803.json

@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "Amazon S3",
+    "contributor": "",
+    "description": "Fix bug in S3 Multipart uploads with FileAsyncRequestBody - ensure that concurrency is limited correctly by bufferSizeInBytes"
+}

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/async/FileAsyncRequestBodySplitHelper.java

@@ -18,7 +18,10 @@
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
 import java.nio.file.attribute.FileTime;
+import java.util.Collections;
+import java.util.HashSet;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
@@ -52,7 +55,7 @@ public final class FileAsyncRequestBodySplitHelper {
 
     private volatile boolean isDone = false;
 
-    private AtomicInteger numAsyncRequestBodiesInFlight = new AtomicInteger(0);
+    private Set<Long> requestBodyStartPositionsInFlight = Collections.synchronizedSet(new HashSet<>());
     private AtomicInteger chunkIndex = new AtomicInteger(0);
     private final FileTime modifiedTimeAtStart;
     private final long sizeAtStart;
@@ -106,9 +109,10 @@ private void sendAsyncRequestBody(SimplePublisher<AsyncRequestBody> simplePublis
 
     private void doSendAsyncRequestBody(SimplePublisher<AsyncRequestBody> simplePublisher) {
         while (shouldSendMore()) {
-            AsyncRequestBody currentAsyncRequestBody = newFileAsyncRequestBody(simplePublisher);
+            long position = chunkSize * chunkIndex.getAndIncrement();
+            AsyncRequestBody currentAsyncRequestBody = newFileAsyncRequestBody(position, simplePublisher);
             simplePublisher.send(currentAsyncRequestBody);
-            numAsyncRequestBodiesInFlight.incrementAndGet();
+            requestBodyStartPositionsInFlight.add(position);
             checkCompletion(simplePublisher, currentAsyncRequestBody);
         }
     }
@@ -126,13 +130,12 @@ private void checkCompletion(SimplePublisher<AsyncRequestBody> simplePublisher,
         }
     }
 
-    private void startNextRequestBody(SimplePublisher<AsyncRequestBody> simplePublisher) {
-        numAsyncRequestBodiesInFlight.decrementAndGet();
+    private void startNextRequestBody(SimplePublisher<AsyncRequestBody> simplePublisher, long completedPosition) {
+        requestBodyStartPositionsInFlight.remove(completedPosition);
         sendAsyncRequestBody(simplePublisher);
     }
 
-    private AsyncRequestBody newFileAsyncRequestBody(SimplePublisher<AsyncRequestBody> simplePublisher) {
-        long position = chunkSize * chunkIndex.getAndIncrement();
+    private AsyncRequestBody newFileAsyncRequestBody(long position, SimplePublisher<AsyncRequestBody> simplePublisher) {
         long numBytesToReadForThisChunk = Math.min(totalContentLength - position, chunkSize);
         FileAsyncRequestBody fileAsyncRequestBody = FileAsyncRequestBody.builder()
                                                                         .path(path)
@@ -142,7 +145,7 @@ private AsyncRequestBody newFileAsyncRequestBody(SimplePublisher<AsyncRequestBod
                                                                         .modifiedTimeAtStart(modifiedTimeAtStart)
                                                                         .sizeAtStart(sizeAtStart)
                                                                         .build();
-        return new FileAsyncRequestBodyWrapper(fileAsyncRequestBody, simplePublisher);
+        return new FileAsyncRequestBodyWrapper(fileAsyncRequestBody, simplePublisher, position);
     }
 
     /**
@@ -153,35 +156,39 @@ private boolean shouldSendMore() {
             return false;
         }
 
-        long currentUsedBuffer = (long) numAsyncRequestBodiesInFlight.get() * bufferPerAsyncRequestBody;
+        long currentUsedBuffer = (long) requestBodyStartPositionsInFlight.size() * bufferPerAsyncRequestBody;
         return currentUsedBuffer + bufferPerAsyncRequestBody <= totalBufferSize;
     }
 
     @SdkTestInternalApi
-    AtomicInteger numAsyncRequestBodiesInFlight() {
-        return numAsyncRequestBodiesInFlight;
+    int numAsyncRequestBodiesInFlight() {
+        return requestBodyStartPositionsInFlight.size();
     }
 
     private final class FileAsyncRequestBodyWrapper implements AsyncRequestBody {
 
         private final FileAsyncRequestBody fileAsyncRequestBody;
         private final SimplePublisher<AsyncRequestBody> simplePublisher;
+        private final long position;
 
         FileAsyncRequestBodyWrapper(FileAsyncRequestBody fileAsyncRequestBody,
-                                    SimplePublisher<AsyncRequestBody> simplePublisher) {
+                                    SimplePublisher<AsyncRequestBody> simplePublisher, long position) {
             this.fileAsyncRequestBody = fileAsyncRequestBody;
             this.simplePublisher = simplePublisher;
+            this.position = position;
         }
 
         @Override
         public void subscribe(Subscriber<? super ByteBuffer> s) {
-            fileAsyncRequestBody.doAfterOnComplete(() -> startNextRequestBody(simplePublisher))
+            fileAsyncRequestBody.doAfterOnComplete(() -> startNextRequestBody(simplePublisher, position))
                                 // The reason we still need to call startNextRequestBody when the subscription is
                                 // cancelled is that upstream could cancel the subscription even though the stream has
                                 // finished successfully before onComplete. If this happens, doAfterOnComplete callback
                                 // will never be invoked, and if the current buffer is full, the publisher will stop
                                 // sending new FileAsyncRequestBody, leading to uncompleted future.
-                                .doAfterOnCancel(() -> startNextRequestBody(simplePublisher))
+                                .doAfterOnCancel(() -> {
+                                    startNextRequestBody(simplePublisher, position);
+                                })
                                 .subscribe(s);
         }
 

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/FileAsyncRequestBodySplitHelperTest.java

@@ -87,12 +87,13 @@ public void split_differentChunkSize_shouldSplitCorrectly(int chunkSize) throws
 
     private static Runnable verifyConcurrentRequests(FileAsyncRequestBodySplitHelper helper, AtomicInteger maxConcurrency) {
         return () -> {
-            int concurrency = helper.numAsyncRequestBodiesInFlight().get();
+            int concurrency = helper.numAsyncRequestBodiesInFlight();
 
             if (concurrency > maxConcurrency.get()) {
                 maxConcurrency.set(concurrency);
             }
-            assertThat(helper.numAsyncRequestBodiesInFlight()).hasValueLessThan(10);
+            assertThat(concurrency).isLessThan(10);
+            assertThat(concurrency).isGreaterThanOrEqualTo(0);
         };
     }
 }

core/sdk-core/src/test/java/software/amazon/awssdk/core/internal/async/SplittingPublisherTestUtils.java

@@ -33,6 +33,7 @@ public static void verifyIndividualAsyncRequestBody(SdkPublisher<AsyncRequestBod
                                                         Path file,
                                                         int chunkSize) throws Exception {
 
+        long contentLength = file.toFile().length();
         List<CompletableFuture<ByteBuffer>> futures = new ArrayList<>();
         publisher.subscribe(requestBody -> {
             CompletableFuture<ByteBuffer> baosFuture = new CompletableFuture<>();
@@ -47,7 +48,6 @@ public static void verifyIndividualAsyncRequestBody(SdkPublisher<AsyncRequestBod
             futures.add(baosFuture);
         }).get(5, TimeUnit.SECONDS);
 
-        long contentLength = file.toFile().length();
         Assertions.assertThat(futures.size()).isEqualTo((int) Math.ceil(contentLength / (double) chunkSize));
 
         for (int i = 0; i < futures.size(); i++) {