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,9 @@ 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,37 @@ 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);
 
-        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 +157,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/CrtS3TransferManager.java

@@ -20,7 +20,8 @@
 import static software.amazon.awssdk.services.s3.crt.S3CrtSdkHttpExecutionAttribute.METAREQUEST_PAUSE_OBSERVABLE;
 import static software.amazon.awssdk.services.s3.internal.crt.S3InternalSdkHttpExecutionAttribute.CRT_PAUSE_RESUME_TOKEN;
 import static software.amazon.awssdk.transfer.s3.internal.GenericS3TransferManager.assertNotUnsupportedArn;
-import static software.amazon.awssdk.transfer.s3.internal.utils.FileUtils.fileNotModified;
+import static software.amazon.awssdk.transfer.s3.internal.TransferManagerHelper.fileModified;
+import static software.amazon.awssdk.transfer.s3.internal.TransferManagerHelper.hasResumeToken;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
@@ -31,7 +32,6 @@
 import software.amazon.awssdk.http.SdkHttpExecutionAttributes;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.internal.crt.S3MetaRequestPauseObservable;
-import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
 import software.amazon.awssdk.transfer.s3.S3TransferManager;
@@ -99,67 +99,15 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
         return new CrtFileUpload(returnFuture, progressUpdater.progress(), observable, uploadFileRequest);
     }
 
-    private FileUpload uploadFromBeginning(ResumableFileUpload resumableFileUpload, boolean fileModified,
-                                           boolean noResumeToken) {
-        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;
-                                                });
-                               });
-        }
-
-        if (noResumeToken) {
-            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 uploadFile(uploadFileRequest);
-    }
-
     @Override
     public FileUpload resumeUploadFile(ResumableFileUpload resumableFileUpload) {
         Validate.paramNotNull(resumableFileUpload, "resumableFileUpload");
 
-        boolean fileModified = !fileNotModified(resumableFileUpload.fileLength(),
-                                                resumableFileUpload.fileLastModified(),
-                                                resumableFileUpload.uploadFileRequest().source());
-
+        boolean fileModified = fileModified(resumableFileUpload, s3AsyncClient);
         boolean noResumeToken = !hasResumeToken(resumableFileUpload);
 
         if (fileModified || noResumeToken) {
-            return uploadFromBeginning(resumableFileUpload, fileModified, noResumeToken);
+            return uploadFile(resumableFileUpload.uploadFileRequest());
         }
 
         return doResumeUpload(resumableFileUpload);
@@ -188,10 +136,6 @@ private static ResumeToken crtResumeToken(ResumableFileUpload resumableFileUploa
                                    .withUploadId(resumableFileUpload.multipartUploadId().orElse(null)));
     }
 
-    private boolean hasResumeToken(ResumableFileUpload resumableFileUpload) {
-        return resumableFileUpload.totalParts().isPresent() && resumableFileUpload.partSizeInBytes().isPresent();
-    }
-
     private PutObjectRequest attachSdkAttribute(PutObjectRequest putObjectRequest,
                                                 Consumer<SdkHttpExecutionAttributes.Builder> builderMutation) {
         SdkHttpExecutionAttributes modifiedAttributes =

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

@@ -15,21 +15,30 @@
 
 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.TransferManagerHelper.fileModified;
+import static software.amazon.awssdk.transfer.s3.internal.TransferManagerHelper.hasResumeToken;
 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.MultipartS3AsyncClient;
+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;
@@ -65,6 +74,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;
@@ -157,6 +167,15 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
                                 .build();
 
         PutObjectRequest putObjectRequest = uploadFileRequest.putObjectRequest();
+        PauseObservable pauseObservable;
+        if (isS3ClientMultipartEnabled()) {
+            pauseObservable = new PauseObservable();
+            Consumer<AwsRequestOverrideConfiguration.Builder> attachPauseObservable =
+                b -> b.putExecutionAttribute(PAUSE_OBSERVABLE, pauseObservable);
+            putObjectRequest = attachSdkAttribute(uploadFileRequest.putObjectRequest(), attachPauseObservable);
+        } else {
+            pauseObservable = null;
+        }
 
         CompletableFuture<CompletedFileUpload> returnFuture = new CompletableFuture<>();
 
@@ -182,8 +201,61 @@ 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 = fileModified(resumableFileUpload, s3AsyncClient);
+        boolean noResumeToken = !hasResumeToken(resumableFileUpload);
+
+        if (fileModified || noResumeToken) {
+            return uploadFile(resumableFileUpload.uploadFileRequest());
+        }
+
+        return doResumeUpload(resumableFileUpload);
+    }
+
+    private boolean isS3ClientMultipartEnabled() {
+        return s3AsyncClient instanceof MultipartS3AsyncClient;
+    }
+
+    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());
+    }
 
-        return new DefaultFileUpload(returnFuture, progressUpdater.progress(), uploadFileRequest);
+    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();
     }
 
     @Override