Merge pull request #104 from olegmoz/87-s3-multipart-abort

#87 - Aborting multipart upload if failed to upload content

Author: g4s8

src/main/java/com/artipie/asto/s3/MultipartUpload.java

@@ -38,6 +38,7 @@
 import org.reactivestreams.Publisher;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
@@ -151,6 +152,21 @@ public CompletionStage<Void> complete() {
         ).thenApply(ignored -> null);
     }
 
+    /**
+     * Aborts the upload.
+     *
+     * @return Completion stage which is completed when success response received from S3.
+     */
+    public CompletionStage<Void> abort() {
+        return this.client.abortMultipartUpload(
+            AbortMultipartUploadRequest.builder()
+                .bucket(this.bucket)
+                .key(this.key.string())
+                .uploadId(this.id)
+                .build()
+        ).thenApply(ignored -> null);
+    }
+
     /**
      * Uploads part.
      *

src/main/java/com/artipie/asto/s3/S3Storage.java

@@ -32,6 +32,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import java.util.stream.Collectors;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.async.SdkPublisher;
@@ -50,6 +51,12 @@
  * Storage that holds data in S3 storage.
  *
  * @since 0.15
+ * @todo #87:60min Do not await abort to complete if save() failed.
+ *  In case uploading content fails inside {@link S3Storage#save(Key, Content)} method
+ *  we are doing abort() for multipart upload.
+ *  Also whole operation does not complete until abort() is complete.
+ *  It would be better to finish save() operation right away and do abort() in background,
+ *  but it makes testing the method difficult.
  */
 public final class S3Storage implements Storage {
 
@@ -123,9 +130,22 @@ public CompletableFuture<Void> save(final Key key, final Content content) {
         ).thenApply(
             created -> new MultipartUpload(this.client, this.bucket, key, created.uploadId())
         ).thenCompose(
-            upload -> upload.upload(content)
-                .thenCompose(ignored -> upload.complete())
-        ).thenApply(response -> null);
+            upload -> upload.upload(content).handle(
+                (ignored, throwable) -> {
+                    final CompletionStage<Void> finished;
+                    if (throwable == null) {
+                        finished = upload.complete();
+                    } else {
+                        final CompletableFuture<Void> promise = new CompletableFuture<>();
+                        finished = promise;
+                        upload.abort().whenComplete(
+                            (result, ex) -> promise.completeExceptionally(throwable)
+                        );
+                    }
+                    return finished;
+                }
+            ).thenCompose(self -> self)
+        );
     }
 
     @Override

src/test/java/com/artipie/asto/S3StorageTest.java

@@ -25,20 +25,26 @@
 
 import com.adobe.testing.s3mock.junit5.S3MockExtension;
 import com.amazonaws.services.s3.AmazonS3;
+import com.amazonaws.services.s3.model.ListMultipartUploadsRequest;
+import com.amazonaws.services.s3.model.MultipartUpload;
 import com.amazonaws.services.s3.model.ObjectMetadata;
 import com.amazonaws.services.s3.model.S3Object;
 import com.artipie.asto.blocking.BlockingStorage;
 import com.artipie.asto.s3.S3Storage;
 import com.google.common.io.ByteStreams;
+import io.reactivex.Flowable;
 import java.io.ByteArrayInputStream;
 import java.net.URI;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import org.hamcrest.MatcherAssert;
 import org.hamcrest.Matchers;
+import org.hamcrest.collection.IsEmptyIterable;
 import org.hamcrest.core.IsEqual;
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.RegisterExtension;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
@@ -50,6 +56,7 @@
  * Tests for {@link S3Storage}.
  *
  * @since 0.15
+ * @checkstyle ClassDataAbstractionCouplingCheck (2 lines)
  */
 class S3StorageTest {
 
@@ -75,6 +82,35 @@ void shouldUploadObjectWhenSave(final AmazonS3 client) throws Exception {
         MatcherAssert.assertThat(downloaded, Matchers.equalTo(data));
     }
 
+    @Test
+    void shouldFailToSaveMultipartUploadWhenFailedToReadContent(final AmazonS3 client) {
+        final String bucket = UUID.randomUUID().toString();
+        client.createBucket(bucket);
+        final String key = "fail";
+        Assertions.assertThrows(
+            Exception.class,
+            () -> this.storage(bucket).save(
+                new Key.From(key),
+                new Content.From(Flowable.error(new IllegalStateException()))
+            ).join()
+        );
+    }
+
+    @Test
+    void shouldAbortMultipartUploadWhenFailedToReadContent(final AmazonS3 client) {
+        final String bucket = UUID.randomUUID().toString();
+        client.createBucket(bucket);
+        final String key = "abort";
+        this.storage(bucket).save(
+            new Key.From(key),
+            new Content.From(Flowable.error(new IllegalStateException()))
+        ).exceptionally(ignore -> null).join();
+        final List<MultipartUpload> uploads = client.listMultipartUploads(
+            new ListMultipartUploadsRequest(bucket)
+        ).getMultipartUploads();
+        MatcherAssert.assertThat(uploads, new IsEmptyIterable<>());
+    }
+
     @Test
     void shouldExistForSavedObject(final AmazonS3 client) {
         final String bucket = UUID.randomUUID().toString();