Fixed the issue where CRT-based S3 Transfer Manager still uses Java to read file in uploadDirectory (#5012)

zoewangg · web-flow · commit e7bcd36cbf7e · 2024-03-12T15:36:04.000-07:00
* AWS-CRT based S3 Transfer Manager now relies on CRT to perform file reading for upload directory.

* Address sonarcloud issues
diff --git a/.changes/next-release/bugfix-S3TransferManager-e2cd1d0.json b/.changes/next-release/bugfix-S3TransferManager-e2cd1d0.json
@@ -0,0 +1,6 @@
+{
+    "type": "bugfix",
+    "category": "S3 Transfer Manager",
+    "contributor": "",
+    "description": "AWS-CRT based S3 Transfer Manager now relies on CRT to perform file reading for upload directory. Related to [#4999](https://github.com/aws/aws-sdk-java-v2/issues/4999)"
+}
diff --git a/services-custom/s3-transfer-manager/pom.xml b/services-custom/s3-transfer-manager/pom.xml
@@ -221,6 +221,11 @@
             <groupId>com.github.tomakehurst</groupId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-junit-jupiter</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>
diff --git a/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/CrtS3TransferManager.java b/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/CrtS3TransferManager.java
@@ -19,7 +19,6 @@
 import static software.amazon.awssdk.services.s3.crt.S3CrtSdkHttpExecutionAttribute.CRT_PROGRESS_LISTENER;
 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 java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
@@ -40,21 +39,18 @@
 import software.amazon.awssdk.transfer.s3.model.ResumableFileUpload;
 import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
-import software.amazon.awssdk.utils.Logger;
 import software.amazon.awssdk.utils.Validate;
 
 /**
  * An implementation of {@link S3TransferManager} that uses CRT-based S3 client under the hood.
  */
 @SdkInternalApi
-class CrtS3TransferManager extends DelegatingS3TransferManager {
-    private static final Logger log = Logger.loggerFor(S3TransferManager.class);
-    private static final PauseResumeHelper PAUSE_RESUME_HELPER = new PauseResumeHelper();
+class CrtS3TransferManager extends GenericS3TransferManager {
     private final S3AsyncClient s3AsyncClient;
 
     CrtS3TransferManager(TransferManagerConfiguration transferConfiguration, S3AsyncClient s3AsyncClient,
                          boolean isDefaultS3AsyncClient) {
-        super(new GenericS3TransferManager(transferConfiguration, s3AsyncClient, isDefaultS3AsyncClient));
+        super(transferConfiguration, s3AsyncClient, isDefaultS3AsyncClient);
         this.s3AsyncClient = s3AsyncClient;
     }
 
@@ -70,7 +66,7 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
             b -> b.put(METAREQUEST_PAUSE_OBSERVABLE, observable)
                   .put(CRT_PROGRESS_LISTENER, progressUpdater.crtProgressListener());
 
-        PutObjectRequest putObjectRequest = attachSdkAttribute(uploadFileRequest.putObjectRequest(), attachObservable);
+        PutObjectRequest putObjectRequest = attachCrtSdkAttribute(uploadFileRequest.putObjectRequest(), attachObservable);
 
         CompletableFuture<CompletedFileUpload> returnFuture = new CompletableFuture<>();
 
@@ -99,28 +95,15 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
     }
 
     @Override
-    public FileUpload resumeUploadFile(ResumableFileUpload resumableFileUpload) {
-        Validate.paramNotNull(resumableFileUpload, "resumableFileUpload");
-
-        boolean fileModified = PAUSE_RESUME_HELPER.fileModified(resumableFileUpload, s3AsyncClient);
-        boolean noResumeToken = !PAUSE_RESUME_HELPER.hasResumeToken(resumableFileUpload);
-
-        if (fileModified || noResumeToken) {
-            return uploadFile(resumableFileUpload.uploadFileRequest());
-        }
-
-        return doResumeUpload(resumableFileUpload);
-    }
-
-    private FileUpload doResumeUpload(ResumableFileUpload resumableFileUpload) {
+    FileUpload doResumeUpload(ResumableFileUpload resumableFileUpload) {
         UploadFileRequest uploadFileRequest = resumableFileUpload.uploadFileRequest();
         PutObjectRequest putObjectRequest = uploadFileRequest.putObjectRequest();
         ResumeToken resumeToken = crtResumeToken(resumableFileUpload);
 
         Consumer<SdkHttpExecutionAttributes.Builder> attachResumeToken =
             b -> b.put(CRT_PAUSE_RESUME_TOKEN, resumeToken);
 
-        PutObjectRequest modifiedPutObjectRequest = attachSdkAttribute(putObjectRequest, attachResumeToken);
+        PutObjectRequest modifiedPutObjectRequest = attachCrtSdkAttribute(putObjectRequest, attachResumeToken);
 
         return uploadFile(uploadFileRequest.toBuilder()
                                            .putObjectRequest(modifiedPutObjectRequest)
@@ -135,8 +118,8 @@ private static ResumeToken crtResumeToken(ResumableFileUpload resumableFileUploa
                                    .withUploadId(resumableFileUpload.multipartUploadId().orElse(null)));
     }
 
-    private PutObjectRequest attachSdkAttribute(PutObjectRequest putObjectRequest,
-                                                Consumer<SdkHttpExecutionAttributes.Builder> builderMutation) {
+    private PutObjectRequest attachCrtSdkAttribute(PutObjectRequest putObjectRequest,
+                                                   Consumer<SdkHttpExecutionAttributes.Builder> builderMutation) {
         SdkHttpExecutionAttributes modifiedAttributes =
             putObjectRequest.overrideConfiguration().map(o -> o.executionAttributes().getAttribute(SDK_HTTP_EXECUTION_ATTRIBUTES))
                             .map(b -> b.toBuilder().applyMutation(builderMutation).build())
diff --git a/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/GenericS3TransferManager.java b/services-custom/s3-transfer-manager/src/main/java/software/amazon/awssdk/transfer/s3/internal/GenericS3TransferManager.java
@@ -124,7 +124,7 @@ class GenericS3TransferManager implements S3TransferManager {
     }
 
     @Override
-    public Upload upload(UploadRequest uploadRequest) {
+    public final Upload upload(UploadRequest uploadRequest) {
         Validate.paramNotNull(uploadRequest, "uploadRequest");
 
         AsyncRequestBody requestBody = uploadRequest.requestBody();
@@ -164,6 +164,9 @@ public Upload upload(UploadRequest uploadRequest) {
         return new DefaultUpload(returnFuture, progressUpdater.progress());
     }
 
+    /**
+     * May be overridden by subclasses to provide customized behavior
+     */
     @Override
     public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
         Validate.paramNotNull(uploadFileRequest, "uploadFileRequest");
@@ -215,7 +218,7 @@ public FileUpload uploadFile(UploadFileRequest uploadFileRequest) {
     }
 
     @Override
-    public FileUpload resumeUploadFile(ResumableFileUpload resumableFileUpload) {
+    public final FileUpload resumeUploadFile(ResumableFileUpload resumableFileUpload) {
         Validate.paramNotNull(resumableFileUpload, "resumableFileUpload");
 
         boolean fileModified = PAUSE_RESUME_HELPER.fileModified(resumableFileUpload, s3AsyncClient);
@@ -233,7 +236,11 @@ private boolean isS3ClientMultipartEnabled() {
         return s3AsyncClient instanceof MultipartS3AsyncClient;
     }
 
-    private FileUpload doResumeUpload(ResumableFileUpload resumableFileUpload) {
+
+    /**
+     * Can be overridden by subclasses to provide different implementation
+     */
+    FileUpload doResumeUpload(ResumableFileUpload resumableFileUpload) {
         UploadFileRequest uploadFileRequest = resumableFileUpload.uploadFileRequest();
         PutObjectRequest putObjectRequest = uploadFileRequest.putObjectRequest();
         S3ResumeToken s3ResumeToken = s3ResumeToken(resumableFileUpload);
@@ -280,7 +287,7 @@ private PutObjectRequest attachSdkAttribute(PutObjectRequest putObjectRequest,
     }
 
     @Override
-    public DirectoryUpload uploadDirectory(UploadDirectoryRequest uploadDirectoryRequest) {
+    public final DirectoryUpload uploadDirectory(UploadDirectoryRequest uploadDirectoryRequest) {
         Validate.paramNotNull(uploadDirectoryRequest, "uploadDirectoryRequest");
 
         try {
@@ -293,7 +300,7 @@ public DirectoryUpload uploadDirectory(UploadDirectoryRequest uploadDirectoryReq
     }
 
     @Override
-    public <ResultT> Download<ResultT> download(DownloadRequest<ResultT> downloadRequest) {
+    public final <ResultT> Download<ResultT> download(DownloadRequest<ResultT> downloadRequest) {
         Validate.paramNotNull(downloadRequest, "downloadRequest");
 
         AsyncResponseTransformer<GetObjectResponse, ResultT> responseTransformer =
@@ -326,7 +333,7 @@ public <ResultT> Download<ResultT> download(DownloadRequest<ResultT> downloadReq
     }
 
     @Override
-    public FileDownload downloadFile(DownloadFileRequest downloadRequest) {
+    public final FileDownload downloadFile(DownloadFileRequest downloadRequest) {
         Validate.paramNotNull(downloadRequest, "downloadFileRequest");
 
         AsyncResponseTransformer<GetObjectResponse, GetObjectResponse> responseTransformer =
@@ -367,7 +374,7 @@ private TransferProgressUpdater doDownloadFile(
     }
 
     @Override
-    public FileDownload resumeDownloadFile(ResumableFileDownload resumableFileDownload) {
+    public final FileDownload resumeDownloadFile(ResumableFileDownload resumableFileDownload) {
         Validate.paramNotNull(resumableFileDownload, "resumableFileDownload");
         CompletableFuture<CompletedFileDownload> returnFuture = new CompletableFuture<>();
         DownloadFileRequest originalDownloadRequest = resumableFileDownload.downloadFileRequest();
@@ -432,7 +439,7 @@ private static void handleException(CompletableFuture<CompletedFileDownload> ret
     }
 
     @Override
-    public DirectoryDownload downloadDirectory(DownloadDirectoryRequest downloadDirectoryRequest) {
+    public final DirectoryDownload downloadDirectory(DownloadDirectoryRequest downloadDirectoryRequest) {
         Validate.paramNotNull(downloadDirectoryRequest, "downloadDirectoryRequest");
 
         try {
@@ -445,7 +452,7 @@ public DirectoryDownload downloadDirectory(DownloadDirectoryRequest downloadDire
     }
 
     @Override
-    public Copy copy(CopyRequest copyRequest) {
+    public final Copy copy(CopyRequest copyRequest) {
         Validate.paramNotNull(copyRequest, "copyRequest");
 
         CompletableFuture<CompletedCopy> returnFuture = new CompletableFuture<>();
@@ -476,7 +483,7 @@ public Copy copy(CopyRequest copyRequest) {
     }
 
     @Override
-    public void close() {
+    public final void close() {
         if (isDefaultS3AsyncClient) {
             IoUtils.closeQuietly(s3AsyncClient, log.logger());
         }
diff --git a/services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/CrtS3TransferManagerTest.java b/services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/CrtS3TransferManagerTest.java
@@ -0,0 +1,110 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.awssdk.transfer.s3.internal;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static software.amazon.awssdk.core.interceptor.SdkInternalExecutionAttribute.SDK_HTTP_EXECUTION_ATTRIBUTES;
+import static software.amazon.awssdk.services.s3.crt.S3CrtSdkHttpExecutionAttribute.METAREQUEST_PAUSE_OBSERVABLE;
+
+import com.google.common.jimfs.Jimfs;
+import java.io.IOException;
+import java.nio.file.FileSystem;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.util.concurrent.CompletableFuture;
+import org.apache.commons.lang3.RandomStringUtils;
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import software.amazon.awssdk.http.SdkHttpExecutionAttributes;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
+import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.transfer.s3.model.UploadDirectoryRequest;
+import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
+
+@ExtendWith(MockitoExtension.class)
+public class CrtS3TransferManagerTest {
+
+    @Mock
+    private S3AsyncClient s3AsyncClient;
+
+    private static Path localDirectory;
+    private static FileSystem jimfs;
+    private CrtS3TransferManager transferManager;
+
+    @BeforeAll
+    public static void setUp() throws IOException {
+        jimfs = Jimfs.newFileSystem();
+        localDirectory = jimfs.getPath("test");
+        Files.createDirectory(localDirectory);
+        Files.write(jimfs.getPath("test", "test.txt"), RandomStringUtils.randomAscii(1024).getBytes());
+    }
+
+    @BeforeEach
+    public void setUpPerMethod() {
+        transferManager = new CrtS3TransferManager(TransferManagerConfiguration.builder().build(),
+                                                                        s3AsyncClient, false);
+    }
+
+    @AfterAll
+    public static void tearDown() throws IOException {
+        jimfs.close();
+    }
+
+    @Test
+    void uploadDirectory_shouldUseCrtUploadFile() {
+        when(s3AsyncClient.putObject(any(PutObjectRequest.class), any(Path.class))).thenReturn(CompletableFuture.completedFuture(PutObjectResponse.builder().build()));
+        transferManager.uploadDirectory(UploadDirectoryRequest.builder().bucket("TEST").source(localDirectory).build())
+                       .completionFuture()
+                       .join();
+
+        verifyCrtInRequestAttributes();
+    }
+
+    @Test
+    void uploadFile_shouldUseCrtUploadFile() {
+        when(s3AsyncClient.putObject(any(PutObjectRequest.class), any(Path.class))).thenReturn(CompletableFuture.completedFuture(PutObjectResponse.builder().build()));
+        transferManager.uploadFile(UploadFileRequest.builder()
+                                                    .putObjectRequest(PutObjectRequest.builder().bucket("test").key("test").build())
+                                                    .source(localDirectory.resolve("test.txt"))
+                                                    .build())
+                       .completionFuture()
+                       .join();
+
+        verifyCrtInRequestAttributes();
+    }
+
+    private void verifyCrtInRequestAttributes() {
+        ArgumentCaptor<PutObjectRequest> requestArgumentCaptor = ArgumentCaptor.forClass(PutObjectRequest.class);
+
+        verify(s3AsyncClient).putObject(requestArgumentCaptor.capture(), ArgumentCaptor.forClass(Path.class).capture());
+
+        PutObjectRequest actual = requestArgumentCaptor.getValue();
+        assertThat(actual.overrideConfiguration()).isPresent();
+        SdkHttpExecutionAttributes attribute = actual.overrideConfiguration().get().executionAttributes().getAttribute(SDK_HTTP_EXECUTION_ATTRIBUTES);
+        assertThat(attribute).isNotNull();
+        assertThat(attribute.getAttribute(METAREQUEST_PAUSE_OBSERVABLE)).isNotNull();
+    }
+}
