feat(storage:s3): multi-part upload: upload parts concurrently

Author: jeqo

checkstyle/suppressions.xml

@@ -22,6 +22,7 @@
     <suppress checks="ClassDataAbstractionCoupling" files=".*Test\.java"/>
     <suppress checks="ClassFanOutComplexity" files=".*Test\.java"/>
     <suppress checks="ClassFanOutComplexity" files="RemoteStorageManager.java"/>
+    <suppress checks="ClassDataAbstractionCoupling" files="S3MultiPartOutputStream.java"/>
     <suppress checks="ClassDataAbstractionCoupling" files="S3StorageConfig.java"/>
     <suppress checks="ClassDataAbstractionCoupling" files="RemoteStorageManager.java"/>
 </suppressions>

storage/s3/src/main/java/io/aiven/kafka/tieredstorage/storage/s3/S3MultiPartOutputStream.java

@@ -18,29 +18,32 @@
 
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
-import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
 import java.util.ArrayList;
-import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.model.AbortMultipartUploadRequest;
 import com.amazonaws.services.s3.model.CompleteMultipartUploadRequest;
-import com.amazonaws.services.s3.model.CompleteMultipartUploadResult;
 import com.amazonaws.services.s3.model.InitiateMultipartUploadRequest;
 import com.amazonaws.services.s3.model.InitiateMultipartUploadResult;
 import com.amazonaws.services.s3.model.PartETag;
 import com.amazonaws.services.s3.model.UploadPartRequest;
-import com.amazonaws.services.s3.model.UploadPartResult;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 /**
  * S3 multipart output stream.
  * Enable uploads to S3 with unknown size by feeding input bytes to multiple parts and upload them on close.
  *
+ * <p>OutputStream is used to write sequentially, but
+ * uploading parts happen asynchronously to reduce full upload latency.
+ * Concurrency happens within the output stream implementation and does not require changes on the callers.
+ *
  * <p>Requires S3 client and starts a multipart transaction when instantiated. Do not reuse.
  */
 public class S3MultiPartOutputStream extends OutputStream {
@@ -54,7 +57,10 @@ public class S3MultiPartOutputStream extends OutputStream {
     final int partSize;
 
     private final String uploadId;
-    private final List<PartETag> partETags = new ArrayList<>();
+    // holds async part upload operations building a list of partETags required when committing
+    private CompletableFuture<ConcurrentLinkedQueue<PartETag>> partUploads =
+        CompletableFuture.completedFuture(new ConcurrentLinkedQueue<>());
+    private final AtomicInteger partNumber = new AtomicInteger(0);
 
     private boolean closed;
 
@@ -88,32 +94,42 @@ public void write(final byte[] b, final int off, final int len) throws IOExcepti
         }
         final ByteBuffer source = ByteBuffer.wrap(b, off, len);
         while (source.hasRemaining()) {
-            final int transferred = Math.min(partBuffer.remaining(), source.remaining());
-            final int offset = source.arrayOffset() + source.position();
-            // TODO: get rid of this array copying
-            partBuffer.put(source.array(), offset, transferred);
-            source.position(source.position() + transferred);
+            final int toCopy = Math.min(partBuffer.remaining(), source.remaining());
+            final int positionAfterCopying = source.position() + toCopy;
+            source.limit(positionAfterCopying);
+            partBuffer.put(source.slice());
+            source.clear(); // reset limit
+            source.position(positionAfterCopying);
             if (!partBuffer.hasRemaining()) {
-                flushBuffer(0, partSize);
+                partBuffer.position(0);
+                partBuffer.limit(partSize);
+                flushBuffer(partBuffer.slice(), partSize);
+                partBuffer.clear();
             }
         }
     }
 
     @Override
     public void close() throws IOException {
         if (partBuffer.position() > 0) {
-            flushBuffer(partBuffer.arrayOffset(), partBuffer.position());
+            final int actualPartSize = partBuffer.position();
+            partBuffer.position(0);
+            partBuffer.limit(actualPartSize);
+            flushBuffer(partBuffer.slice(), actualPartSize);
         }
         if (Objects.nonNull(uploadId)) {
-            if (!partETags.isEmpty()) {
+            if (partNumber.get() > 0) {
                 try {
-                    final CompleteMultipartUploadRequest request =
-                        new CompleteMultipartUploadRequest(bucketName, key, uploadId, partETags);
-                    final CompleteMultipartUploadResult result = client.completeMultipartUpload(request);
+                    // wait for all uploads to complete successfully before committing
+                    final ConcurrentLinkedQueue<PartETag> tagsQueue = partUploads.get(); // TODO: maybe set a timeout?
+                    final ArrayList<PartETag> partETags = new ArrayList<>(tagsQueue);
+                    final var request = new CompleteMultipartUploadRequest(bucketName, key, uploadId, partETags);
+                    final var result = client.completeMultipartUpload(request);
                     log.debug("Completed multipart upload {} with result {}", uploadId, result);
                 } catch (final Exception e) {
                     log.error("Failed to complete multipart upload {}, aborting transaction", uploadId, e);
                     client.abortMultipartUpload(new AbortMultipartUploadRequest(bucketName, key, uploadId));
+                    throw new IOException(e);
                 }
             } else {
                 client.abortMultipartUpload(new AbortMultipartUploadRequest(bucketName, key, uploadId));
@@ -122,31 +138,24 @@ public void close() throws IOException {
         closed = true;
     }
 
-    private void flushBuffer(final int offset,
-                             final int actualPartSize) throws IOException {
-        try {
-            final ByteArrayInputStream in = new ByteArrayInputStream(partBuffer.array(), offset, actualPartSize);
-            uploadPart(in, actualPartSize);
-            partBuffer.clear();
-        } catch (final Exception e) {
-            log.error("Failed to upload part in multipart upload {}, aborting transaction", uploadId, e);
-            client.abortMultipartUpload(new AbortMultipartUploadRequest(bucketName, key, uploadId));
-            closed = true;
-            throw new IOException(e);
-        }
-    }
+    private void flushBuffer(final ByteBuffer partBuffer, final int actualPartSize) {
+        final byte[] partContent = new byte[actualPartSize];
+        partBuffer.get(partContent, 0, actualPartSize);
+
+        final var uploadPartRequest = new UploadPartRequest()
+            .withBucketName(bucketName)
+            .withKey(key)
+            .withUploadId(uploadId)
+            .withPartSize(actualPartSize)
+            .withPartNumber(partNumber.incrementAndGet())
+            .withInputStream(new ByteArrayInputStream(partContent));
 
-    private void uploadPart(final InputStream in, final int actualPartSize) {
-        final int partNumber = partETags.size() + 1;
-        final UploadPartRequest uploadPartRequest =
-            new UploadPartRequest()
-                .withBucketName(bucketName)
-                .withKey(key)
-                .withUploadId(uploadId)
-                .withPartSize(actualPartSize)
-                .withPartNumber(partNumber)
-                .withInputStream(in);
-        final UploadPartResult uploadResult = client.uploadPart(uploadPartRequest);
-        partETags.add(uploadResult.getPartETag());
+        // Run request async
+        partUploads = partUploads.thenCombine(
+            CompletableFuture.supplyAsync(() -> client.uploadPart(uploadPartRequest)),
+            (partETags, result) -> {
+                partETags.add(result.getPartETag());
+                return partETags;
+            });
     }
 }

storage/s3/src/test/java/io/aiven/kafka/tieredstorage/storage/s3/S3MultiPartOutputStreamTest.java

@@ -16,11 +16,12 @@
 
 package io.aiven.kafka.tieredstorage.storage.s3;
 
-import java.io.ByteArrayInputStream;
 import java.io.IOException;
-import java.util.ArrayList;
-import java.util.List;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Random;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 import com.amazonaws.services.s3.AmazonS3;
 import com.amazonaws.services.s3.model.AbortMultipartUploadRequest;
@@ -83,7 +84,7 @@ void sendAbortForAnyExceptionWhileWriting() {
                      new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, 100, mockedS3)) {
                 out.write(new byte[] {1, 2, 3});
             }
-        }).isInstanceOf(IOException.class).hasCause(testException);
+        }).isInstanceOf(IOException.class).hasRootCause(testException);
 
         verify(mockedS3).initiateMultipartUpload(any(InitiateMultipartUploadRequest.class));
         verify(mockedS3).uploadPart(any(UploadPartRequest.class));
@@ -101,7 +102,7 @@ void sendAbortForAnyExceptionWhenClose() throws Exception {
         when(mockedS3.uploadPart(any(UploadPartRequest.class)))
             .thenThrow(RuntimeException.class);
 
-        final S3MultiPartOutputStream out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, 10, mockedS3);
+        final var out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, 10, mockedS3);
 
         final byte[] buffer = new byte[5];
         random.nextBytes(buffer);
@@ -127,110 +128,124 @@ void writesOneByte() throws Exception {
         when(mockedS3.completeMultipartUpload(completeMultipartUploadRequestCaptor.capture()))
             .thenReturn(new CompleteMultipartUploadResult());
 
-        try (final S3MultiPartOutputStream out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, 100, mockedS3)) {
+        try (final var out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, 100, mockedS3)) {
             out.write(1);
         }
 
         verify(mockedS3).initiateMultipartUpload(any(InitiateMultipartUploadRequest.class));
         verify(mockedS3).uploadPart(any(UploadPartRequest.class));
         verify(mockedS3).completeMultipartUpload(any(CompleteMultipartUploadRequest.class));
 
+        final UploadPartRequest value = uploadPartRequestCaptor.getValue();
         assertUploadPartRequest(
-            uploadPartRequestCaptor.getValue(),
+            value,
+            value.getInputStream().readAllBytes(),
             1,
             1,
             new byte[] {1});
         assertCompleteMultipartUploadRequest(
             completeMultipartUploadRequestCaptor.getValue(),
-            List.of(new PartETag(1, "SOME_ETAG"))
+            Map.of(1, "SOME_ETAG")
         );
     }
 
     @Test
     void writesMultipleMessages() throws Exception {
         final int bufferSize = 10;
-        final byte[] message = new byte[bufferSize];
 
         when(mockedS3.initiateMultipartUpload(any(InitiateMultipartUploadRequest.class)))
             .thenReturn(newInitiateMultipartUploadResult());
-        when(mockedS3.uploadPart(uploadPartRequestCaptor.capture()))
-            .thenAnswer(a -> {
-                final UploadPartRequest up = a.getArgument(0);
+
+        // capturing requests and contents from concurrent threads
+        final Map<Integer, UploadPartRequest> uploadPartRequests = new ConcurrentHashMap<>();
+        final Map<Integer, byte[]> uploadPartContents = new ConcurrentHashMap<>();
+        when(mockedS3.uploadPart(any(UploadPartRequest.class)))
+            .thenAnswer(answer -> {
+                final UploadPartRequest up = answer.getArgument(0);
+                //emulate behave of S3 client otherwise we will get wrong array in the memory
+                uploadPartRequests.put(up.getPartNumber(), up);
+                uploadPartContents.put(up.getPartNumber(), up.getInputStream().readAllBytes());
+
                 return newUploadPartResult(up.getPartNumber(), "SOME_TAG#" + up.getPartNumber());
             });
-        when(mockedS3.completeMultipartUpload(completeMultipartUploadRequestCaptor.capture()))
+        when(mockedS3.completeMultipartUpload(any(CompleteMultipartUploadRequest.class)))
             .thenReturn(new CompleteMultipartUploadResult());
 
-        final List<byte[]> expectedMessagesList = new ArrayList<>();
-        try (final S3MultiPartOutputStream out =
-                 new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, bufferSize, mockedS3)) {
+        final Map<Integer, byte[]> expectedMessageParts = new HashMap<>();
+        try (final var out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, bufferSize, mockedS3)) {
             for (int i = 0; i < 3; i++) {
+                final byte[] message = new byte[bufferSize];
                 random.nextBytes(message);
                 out.write(message, 0, message.length);
-                expectedMessagesList.add(message);
+                expectedMessageParts.put(i + 1, message);
             }
         }
 
         verify(mockedS3).initiateMultipartUpload(any(InitiateMultipartUploadRequest.class));
         verify(mockedS3, times(3)).uploadPart(any(UploadPartRequest.class));
         verify(mockedS3).completeMultipartUpload(any(CompleteMultipartUploadRequest.class));
 
-        final List<UploadPartRequest> uploadRequests = uploadPartRequestCaptor.getAllValues();
-        int counter = 0;
-        for (final byte[] expectedMessage : expectedMessagesList) {
+        for (final Integer part : expectedMessageParts.keySet()) {
             assertUploadPartRequest(
-                uploadRequests.get(counter),
+                uploadPartRequests.get(part),
+                uploadPartContents.get(part),
                 bufferSize,
-                counter + 1,
-                expectedMessage);
-            counter++;
+                part,
+                expectedMessageParts.get(part)
+            );
         }
         assertCompleteMultipartUploadRequest(
             completeMultipartUploadRequestCaptor.getValue(),
-            List.of(new PartETag(1, "SOME_TAG#1"),
-                new PartETag(2, "SOME_TAG#2"),
-                new PartETag(3, "SOME_TAG#3"))
+            Map.of(1, "SOME_TAG#1",
+                2, "SOME_TAG#2",
+                3, "SOME_TAG#3")
         );
     }
 
     @Test
     void writesTailMessages() throws Exception {
         final int messageSize = 20;
 
-        final List<UploadPartRequest> uploadPartRequests = new ArrayList<>();
-
         when(mockedS3.initiateMultipartUpload(any(InitiateMultipartUploadRequest.class)))
             .thenReturn(newInitiateMultipartUploadResult());
+
+        // capturing requests and contents from concurrent threads
+        final Map<Integer, UploadPartRequest> uploadPartRequests = new ConcurrentHashMap<>();
+        final Map<Integer, byte[]> uploadPartContents = new ConcurrentHashMap<>();
         when(mockedS3.uploadPart(any(UploadPartRequest.class)))
-            .thenAnswer(a -> {
-                final UploadPartRequest up = a.getArgument(0);
+            .thenAnswer(answer -> {
+                final UploadPartRequest up = answer.getArgument(0);
                 //emulate behave of S3 client otherwise we will get wrong array in the memory
-                up.setInputStream(new ByteArrayInputStream(up.getInputStream().readAllBytes()));
-                uploadPartRequests.add(up);
+                uploadPartRequests.put(up.getPartNumber(), up);
+                uploadPartContents.put(up.getPartNumber(), up.getInputStream().readAllBytes());
 
                 return newUploadPartResult(up.getPartNumber(), "SOME_TAG#" + up.getPartNumber());
             });
-        when(mockedS3.completeMultipartUpload(completeMultipartUploadRequestCaptor.capture()))
+        when(mockedS3.completeMultipartUpload(any(CompleteMultipartUploadRequest.class)))
             .thenReturn(new CompleteMultipartUploadResult());
 
-        final byte[] message = new byte[messageSize];
-
         final byte[] expectedFullMessage = new byte[messageSize + 10];
         final byte[] expectedTailMessage = new byte[10];
 
-        final S3MultiPartOutputStream
-            out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, messageSize + 10, mockedS3);
-        random.nextBytes(message);
-        out.write(message);
-        System.arraycopy(message, 0, expectedFullMessage, 0, message.length);
-        random.nextBytes(message);
-        out.write(message);
-        System.arraycopy(message, 0, expectedFullMessage, 20, 10);
-        System.arraycopy(message, 10, expectedTailMessage, 0, 10);
+        final var out = new S3MultiPartOutputStream(BUCKET_NAME, FILE_KEY, messageSize + 10, mockedS3);
+        {
+            final byte[] message = new byte[messageSize];
+            random.nextBytes(message);
+            out.write(message);
+            System.arraycopy(message, 0, expectedFullMessage, 0, message.length);
+        }
+        {
+            final byte[] message = new byte[messageSize];
+            random.nextBytes(message);
+            out.write(message);
+            System.arraycopy(message, 0, expectedFullMessage, 20, 10);
+            System.arraycopy(message, 10, expectedTailMessage, 0, 10);
+        }
         out.close();
 
-        assertUploadPartRequest(uploadPartRequests.get(0), 30, 1, expectedFullMessage);
-        assertUploadPartRequest(uploadPartRequests.get(1), 10, 2, expectedTailMessage);
+        assertThat(uploadPartRequests).hasSize(2);
+        assertUploadPartRequest(uploadPartRequests.get(1), uploadPartContents.get(1), 30, 1, expectedFullMessage);
+        assertUploadPartRequest(uploadPartRequests.get(2), uploadPartContents.get(2), 10, 2, expectedTailMessage);
 
         verify(mockedS3).initiateMultipartUpload(any(InitiateMultipartUploadRequest.class));
         verify(mockedS3, times(2)).uploadPart(any(UploadPartRequest.class));
@@ -251,6 +266,7 @@ private static UploadPartResult newUploadPartResult(final int partNumber, final
     }
 
     private static void assertUploadPartRequest(final UploadPartRequest uploadPartRequest,
+                                                final byte[] bytes,
                                                 final int expectedPartSize,
                                                 final int expectedPartNumber,
                                                 final byte[] expectedBytes) {
@@ -259,23 +275,17 @@ private static void assertUploadPartRequest(final UploadPartRequest uploadPartRe
         assertThat(uploadPartRequest.getPartNumber()).isEqualTo(expectedPartNumber);
         assertThat(uploadPartRequest.getBucketName()).isEqualTo(BUCKET_NAME);
         assertThat(uploadPartRequest.getKey()).isEqualTo(FILE_KEY);
-        assertThat(uploadPartRequest.getInputStream()).hasBinaryContent(expectedBytes);
+        assertThat(bytes).isEqualTo(expectedBytes);
     }
 
     private static void assertCompleteMultipartUploadRequest(final CompleteMultipartUploadRequest request,
-                                                             final List<PartETag> expectedETags) {
+                                                             final Map<Integer, String> expectedETags) {
         assertThat(request.getBucketName()).isEqualTo(BUCKET_NAME);
         assertThat(request.getKey()).isEqualTo(FILE_KEY);
         assertThat(request.getUploadId()).isEqualTo(UPLOAD_ID);
-        assertThat(request.getPartETags()).hasSameSizeAs(expectedETags);
-
-        for (int i = 0; i < expectedETags.size(); i++) {
-            final PartETag expectedETag = expectedETags.get(i);
-            final PartETag etag = request.getPartETags().get(i);
-
-            assertThat(etag.getPartNumber()).isEqualTo(expectedETag.getPartNumber());
-            assertThat(etag.getETag()).isEqualTo(expectedETag.getETag());
-        }
+        final Map<Integer, String> tags = request.getPartETags().stream()
+            .collect(Collectors.toMap(PartETag::getPartNumber, PartETag::getETag));
+        assertThat(tags).containsExactlyInAnyOrderEntriesOf(expectedETags);
     }
 
     private static void assertAbortMultipartUploadRequest(final AbortMultipartUploadRequest request) {