Pause/Resume Upload for Transfer Manager with Java S3Client

services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3IntegrationTestBase.java

@@ -68,7 +68,7 @@ public static void setUpForAllIntegTests() throws Exception {
         Log.initLoggingToStdout(Log.LogLevel.Warn);
         System.setProperty("aws.crt.debugnative", "true");
         s3 = s3ClientBuilder().build();
-        s3Async = s3AsyncClientBuilder().build();
+        s3Async = s3AsyncClientBuilder().multipartEnabled(true).build();
         s3CrtAsync = S3CrtAsyncClient.builder()
                                      .credentialsProvider(CREDENTIALS_PROVIDER_CHAIN)
                                      .region(DEFAULT_REGION)

services-custom/s3-transfer-manager/src/it/java/software/amazon/awssdk/transfer/s3/S3TransferManagerUploadPauseResumeIntegrationTest.java

@@ -16,7 +16,6 @@
 package software.amazon.awssdk.transfer.s3;
 
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
-import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
 import static software.amazon.awssdk.testutils.service.S3BucketUtils.temporaryBucketName;
 import static software.amazon.awssdk.transfer.s3.SizeConstant.MB;
 
@@ -26,9 +25,12 @@
 import java.time.Duration;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.stream.Stream;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
-import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.MethodSource;
 import software.amazon.awssdk.core.retry.backoff.FixedDelayBackoffStrategy;
 import software.amazon.awssdk.core.waiters.AsyncWaiter;
 import software.amazon.awssdk.core.waiters.Waiter;
@@ -69,30 +71,42 @@ public static void cleanup() {
         executorService.shutdown();
     }
 
-    @Test
-    void pause_singlePart_shouldResume() {
+    enum S3ClientType {
+        CRT, JAVA
+    }
+
+    private static Stream<Arguments> s3ClientType() {
+        return Stream.of(Arguments.of(S3ClientType.CRT), Arguments.of(S3ClientType.JAVA));
+    }
+
+    @ParameterizedTest
+    @MethodSource("s3ClientType")
+    void pause_singlePart_shouldResume(S3ClientType s3ClientType) {
+        S3TransferManager transferManager = s3ClientType == S3ClientType.CRT ? tmCrt : tmJava;
         UploadFileRequest request = UploadFileRequest.builder()
                                                      .putObjectRequest(b -> b.bucket(BUCKET).key(KEY))
                                                      .source(smallFile)
                                                      .build();
-        FileUpload fileUpload = tmCrt.uploadFile(request);
+        FileUpload fileUpload = transferManager.uploadFile(request);
         ResumableFileUpload resumableFileUpload = fileUpload.pause();
         log.debug(() -> "Paused: " + resumableFileUpload);
 
         validateEmptyResumeToken(resumableFileUpload);
 
-        FileUpload resumedUpload = tmCrt.resumeUploadFile(resumableFileUpload);
+        FileUpload resumedUpload = transferManager.resumeUploadFile(resumableFileUpload);
         resumedUpload.completionFuture().join();
     }
 
-    @Test
-    void pause_fileNotChanged_shouldResume() {
+    @ParameterizedTest
+    @MethodSource("s3ClientType")
+    void pause_fileNotChanged_shouldResume(S3ClientType s3ClientType) {
+        S3TransferManager transferManager = s3ClientType == S3ClientType.CRT ? tmCrt : tmJava;
         UploadFileRequest request = UploadFileRequest.builder()
                                                      .putObjectRequest(b -> b.bucket(BUCKET).key(KEY))
                                                      .addTransferListener(LoggingTransferListener.create())
                                                      .source(largeFile)
                                                      .build();
-        FileUpload fileUpload = tmCrt.uploadFile(request);
+        FileUpload fileUpload = transferManager.uploadFile(request);
         waitUntilMultipartUploadExists();
         ResumableFileUpload resumableFileUpload = fileUpload.pause();
         log.debug(() -> "Paused: " + resumableFileUpload);
@@ -103,33 +117,43 @@ void pause_fileNotChanged_shouldResume() {
 
         verifyMultipartUploadIdExists(resumableFileUpload);
 
-        FileUpload resumedUpload = tmCrt.resumeUploadFile(resumableFileUpload);
+        FileUpload resumedUpload = transferManager.resumeUploadFile(resumableFileUpload);
         resumedUpload.completionFuture().join();
     }
 
-    @Test
-    void pauseImmediately_resume_shouldStartFromBeginning() {
+    @ParameterizedTest
+    @MethodSource("s3ClientType")
+    void pauseImmediately_resume_shouldStartFromBeginning(S3ClientType s3ClientType) {
+        S3TransferManager transferManager = s3ClientType == S3ClientType.CRT ? tmCrt : tmJava;
         UploadFileRequest request = UploadFileRequest.builder()
                                                    .putObjectRequest(b -> b.bucket(BUCKET).key(KEY))
                                                    .source(largeFile)
                                                    .build();
-        FileUpload fileUpload = tmCrt.uploadFile(request);
+        FileUpload fileUpload = transferManager.uploadFile(request);
         ResumableFileUpload resumableFileUpload = fileUpload.pause();
         log.debug(() -> "Paused: " + resumableFileUpload);
 
-        validateEmptyResumeToken(resumableFileUpload);
+        if (s3ClientType == S3ClientType.CRT) {
+            validateEmptyResumeToken(resumableFileUpload);
+        } else {
+            // join() is called on CreateMultipartUpload so resume token will be returned
+            assertThat(resumableFileUpload.transferredParts().isPresent()).isTrue();
+            assertThat(resumableFileUpload.transferredParts().getAsLong()).isZero();
+        }
 
-        FileUpload resumedUpload = tmCrt.resumeUploadFile(resumableFileUpload);
+        FileUpload resumedUpload = transferManager.resumeUploadFile(resumableFileUpload);
         resumedUpload.completionFuture().join();
     }
 
-    @Test
-    void pause_fileChanged_resumeShouldStartFromBeginning() throws Exception {
+    @ParameterizedTest
+    @MethodSource("s3ClientType")
+    void pause_fileChanged_resumeShouldStartFromBeginning(S3ClientType s3ClientType) throws Exception {
+        S3TransferManager transferManager = s3ClientType == S3ClientType.CRT ? tmCrt : tmJava;
         UploadFileRequest request = UploadFileRequest.builder()
                                                      .putObjectRequest(b -> b.bucket(BUCKET).key(KEY))
                                                      .source(largeFile)
                                                      .build();
-        FileUpload fileUpload = tmCrt.uploadFile(request);
+        FileUpload fileUpload = transferManager.uploadFile(request);
         waitUntilMultipartUploadExists();
         ResumableFileUpload resumableFileUpload = fileUpload.pause();
         log.debug(() -> "Paused: " + resumableFileUpload);
@@ -139,13 +163,16 @@ void pause_fileChanged_resumeShouldStartFromBeginning() throws Exception {
         assertThat(resumableFileUpload.totalParts()).isNotEmpty();
         verifyMultipartUploadIdExists(resumableFileUpload);
 
+        byte[] originalBytes = Files.readAllBytes(largeFile.toPath());
         byte[] bytes = "helloworld".getBytes(StandardCharsets.UTF_8);
         Files.write(largeFile.toPath(), bytes);
 
-        FileUpload resumedUpload = tmCrt.resumeUploadFile(resumableFileUpload);
+        FileUpload resumedUpload = transferManager.resumeUploadFile(resumableFileUpload);
         resumedUpload.completionFuture().join();
         verifyMultipartUploadIdNotExist(resumableFileUpload);
         assertThat(resumedUpload.progress().snapshot().totalBytes()).hasValue(bytes.length);
+
+        Files.write(largeFile.toPath(), originalBytes);
     }
 
     private void verifyMultipartUploadIdExists(ResumableFileUpload resumableFileUpload) {

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/GenericS3TransferManager.java

@@ -15,24 +15,31 @@
 
 package software.amazon.awssdk.transfer.s3.internal;
 
+import static software.amazon.awssdk.services.s3.internal.multipart.UploadObjectHelper.PAUSE_OBSERVABLE;
+import static software.amazon.awssdk.services.s3.internal.multipart.UploadObjectHelper.RESUME_TOKEN;
 import static software.amazon.awssdk.transfer.s3.SizeConstant.MB;
 import static software.amazon.awssdk.transfer.s3.internal.utils.ResumableRequestConverter.toDownloadFileRequestAndTransformer;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.function.Consumer;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.arns.Arn;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.core.FileTransformerConfiguration;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.core.exception.SdkClientException;
 import software.amazon.awssdk.core.exception.SdkException;
 import software.amazon.awssdk.core.internal.async.FileAsyncRequestBody;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.internal.multipart.PauseObservable;
+import software.amazon.awssdk.services.s3.internal.multipart.S3ResumeToken;
 import software.amazon.awssdk.services.s3.internal.resource.S3AccessPointResource;
 import software.amazon.awssdk.services.s3.internal.resource.S3ArnConverter;
 import software.amazon.awssdk.services.s3.internal.resource.S3Resource;
+import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.CopyObjectResponse;
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
@@ -49,6 +56,7 @@
 import software.amazon.awssdk.transfer.s3.internal.model.DefaultUpload;
 import software.amazon.awssdk.transfer.s3.internal.progress.ResumeTransferProgress;
 import software.amazon.awssdk.transfer.s3.internal.progress.TransferProgressUpdater;
+import software.amazon.awssdk.transfer.s3.internal.utils.FileUtils;
 import software.amazon.awssdk.transfer.s3.model.CompletedCopy;
 import software.amazon.awssdk.transfer.s3.model.CompletedDownload;
 import software.amazon.awssdk.transfer.s3.model.CompletedFileDownload;
@@ -65,6 +73,7 @@
 import software.amazon.awssdk.transfer.s3.model.FileDownload;
 import software.amazon.awssdk.transfer.s3.model.FileUpload;
 import software.amazon.awssdk.transfer.s3.model.ResumableFileDownload;
+import software.amazon.awssdk.transfer.s3.model.ResumableFileUpload;
 import software.amazon.awssdk.transfer.s3.model.Upload;
 import software.amazon.awssdk.transfer.s3.model.UploadDirectoryRequest;
 import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
@@ -156,7 +165,10 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
                                 .chunkSizeInBytes(DEFAULT_FILE_UPLOAD_CHUNK_SIZE)
                                 .build();
 
-        PutObjectRequest putObjectRequest = uploadFileRequest.putObjectRequest();
+        PauseObservable pauseObservable = new PauseObservable();
+        Consumer<AwsRequestOverrideConfiguration.Builder> attachPauseObservable =
+            b -> b.putExecutionAttribute(PAUSE_OBSERVABLE, pauseObservable);
+        PutObjectRequest putObjectRequest = attachSdkAttribute(uploadFileRequest.putObjectRequest(), attachPauseObservable);
 
         CompletableFuture<CompletedFileUpload> returnFuture = new CompletableFuture<>();
 
@@ -182,8 +194,109 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
         } catch (Throwable throwable) {
             returnFuture.completeExceptionally(throwable);
         }
+        return new DefaultFileUpload(returnFuture, progressUpdater.progress(), pauseObservable, uploadFileRequest);
+    }
+
+    @Override
+    public FileUpload resumeUploadFile(ResumableFileUpload resumableFileUpload) {
+        Validate.paramNotNull(resumableFileUpload, "resumableFileUpload");
+
+        boolean fileModified = !FileUtils.fileNotModified(resumableFileUpload.fileLength(),
+                                                          resumableFileUpload.fileLastModified(),
+                                                          resumableFileUpload.uploadFileRequest().source());
+
+        boolean noResumeToken = !hasResumeToken(resumableFileUpload);
+
+        if (fileModified || noResumeToken) {
+            return uploadFromBeginning(resumableFileUpload, fileModified);
+        }
+
+        return doResumeUpload(resumableFileUpload);
+    }
+
+    private FileUpload doResumeUpload(ResumableFileUpload resumableFileUpload) {
+        UploadFileRequest uploadFileRequest = resumableFileUpload.uploadFileRequest();
+        PutObjectRequest putObjectRequest = uploadFileRequest.putObjectRequest();
+        S3ResumeToken s3ResumeToken = s3ResumeToken(resumableFileUpload);
+
+        Consumer<AwsRequestOverrideConfiguration.Builder> attachResumeToken =
+            b -> b.putExecutionAttribute(RESUME_TOKEN, s3ResumeToken);
+
+        PutObjectRequest modifiedPutObjectRequest = attachSdkAttribute(putObjectRequest, attachResumeToken);
+
+        return uploadFile(uploadFileRequest.toBuilder()
+                                           .putObjectRequest(modifiedPutObjectRequest)
+                                           .build());
+    }
+
+    private static S3ResumeToken s3ResumeToken(ResumableFileUpload resumableFileUpload) {
+        return new S3ResumeToken(resumableFileUpload.multipartUploadId().get(),
+                                 resumableFileUpload.partSizeInBytes().getAsLong(),
+                                 resumableFileUpload.totalParts().getAsLong(),
+                                 resumableFileUpload.transferredParts().getAsLong());
+    }
+
+    private boolean hasResumeToken(ResumableFileUpload resumableFileUpload) {
+        return resumableFileUpload.totalParts().isPresent() && resumableFileUpload.partSizeInBytes().isPresent();
+    }
+
+    private PutObjectRequest attachSdkAttribute(PutObjectRequest putObjectRequest,
+                                                Consumer<AwsRequestOverrideConfiguration.Builder> builderMutation) {
+        AwsRequestOverrideConfiguration modifiedRequestOverrideConfig =
+            putObjectRequest.overrideConfiguration()
+                            .map(o -> o.toBuilder().applyMutation(builderMutation).build())
+                            .orElseGet(() -> AwsRequestOverrideConfiguration.builder()
+                                                                            .applyMutation(builderMutation)
+                                                                            .build());
+
+        return putObjectRequest.toBuilder()
+                               .overrideConfiguration(modifiedRequestOverrideConfig)
+                               .build();
+    }
+
+    private FileUpload uploadFromBeginning(ResumableFileUpload resumableFileUpload, boolean fileModified) {
+        UploadFileRequest uploadFileRequest = resumableFileUpload.uploadFileRequest();
+        PutObjectRequest putObjectRequest = uploadFileRequest.putObjectRequest();
+        if (fileModified) {
+            log.debug(() -> String.format("The file (%s) has been modified since "
+                                          + "the last pause. " +
+                                          "The SDK will upload the requested object in bucket"
+                                          + " (%s) with key (%s) from "
+                                          + "the "
+                                          + "beginning.",
+                                          uploadFileRequest.source(),
+                                          putObjectRequest.bucket(),
+                                          putObjectRequest.key()));
+            resumableFileUpload.multipartUploadId()
+                               .ifPresent(id -> {
+                                   log.debug(() -> "Aborting previous upload with multipartUploadId: " + id);
+                                   s3AsyncClient.abortMultipartUpload(
+                                                    AbortMultipartUploadRequest.builder()
+                                                                               .bucket(putObjectRequest.bucket())
+                                                                               .key(putObjectRequest.key())
+                                                                               .uploadId(id)
+                                                                               .build())
+                                                .exceptionally(t -> {
+                                                    log.warn(() -> String.format("Failed to abort previous multipart upload "
+                                                                                 + "(id: %s)"
+                                                                                 + ". You may need to call "
+                                                                                 + "S3AsyncClient#abortMultiPartUpload to "
+                                                                                 + "free all storage consumed by"
+                                                                                 + " all parts. ",
+                                                                                 id), t);
+                                                    return null;
+                                                });
+                               });
+        }
+
+        log.debug(() -> String.format("No resume token is found. " +
+                                      "The SDK will upload the requested object in bucket"
+                                      + " (%s) with key (%s) from "
+                                      + "the beginning.",
+                                      putObjectRequest.bucket(),
+                                      putObjectRequest.key()));
 
-        return new DefaultFileUpload(returnFuture, progressUpdater.progress(), uploadFileRequest);
+        return uploadFile(uploadFileRequest);
     }
 
     @Override

services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/model/CrtFileUpload.java

@@ -148,7 +148,7 @@ public int hashCode() {
 
     @Override
     public String toString() {
-        return ToString.builder("DefaultFileUpload")
+        return ToString.builder("CrtFileUpload")
                        .add("completionFuture", completionFuture)
                        .add("progress", progress)
                        .add("request", request)