Add request cancellation logic and tests

pom.xml

@@ -112,7 +112,7 @@
         <rxjava.version>2.2.21</rxjava.version>
         <commons-codec.verion>1.10</commons-codec.verion>
         <jmh.version>1.29</jmh.version>
-        <awscrt.version>0.12.8</awscrt.version>
+        <awscrt.version>0.13.2</awscrt.version>
 
         <!--Test dependencies -->
         <junit.version>4.13.1</junit.version>

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

@@ -16,9 +16,11 @@
 package software.amazon.awssdk.transfer.s3.internal;
 
 import com.amazonaws.s3.S3NativeClient;
+import com.amazonaws.s3.model.GetObjectOutput;
 import com.amazonaws.s3.model.PutObjectOutput;
 import java.util.concurrent.CompletableFuture;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.annotations.SdkTestInternalApi;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
@@ -28,6 +30,7 @@
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
 
 @SdkInternalApi
 public final class DefaultS3CrtAsyncClient implements S3CrtAsyncClient {
@@ -55,6 +58,13 @@ public DefaultS3CrtAsyncClient(DefaultS3CrtClientBuilder builder) {
                                                  configuration.maxConcurrency());
     }
 
+    @SdkTestInternalApi
+    DefaultS3CrtAsyncClient(S3NativeClientConfiguration configuration,
+                            S3NativeClient nativeClient) {
+        this.configuration = configuration;
+        this.s3NativeClient = nativeClient;
+    }
+
     @Override
     public <ReturnT> CompletableFuture<ReturnT> getObject(
         GetObjectRequest getObjectRequest, AsyncResponseTransformer<GetObjectResponse, ReturnT> asyncResponseTransformer) {
@@ -65,7 +75,8 @@ public <ReturnT> CompletableFuture<ReturnT> getObject(
 
         CompletableFuture<ReturnT> adapterFuture = adapter.transformerFuture();
 
-        s3NativeClient.getObject(crtGetObjectRequest, adapter);
+        CompletableFuture<GetObjectOutput> crtFuture = s3NativeClient.getObject(crtGetObjectRequest, adapter);
+        CompletableFutureUtils.forwardExceptionTo(future, crtFuture);
 
         adapterFuture.whenComplete((r, t) -> {
             if (t == null) {
@@ -76,7 +87,7 @@ public <ReturnT> CompletableFuture<ReturnT> getObject(
             // TODO: Offload to future completion thread
         });
 
-        return future;
+        return CompletableFutureUtils.forwardExceptionTo(future, adapterFuture);
     }
 
     @Override
@@ -94,9 +105,12 @@ public CompletableFuture<PutObjectResponse> putObject(PutObjectRequest putObject
                                                                                                        requestDataSupplier);
 
         CompletableFuture<SdkHttpResponse> httpResponseFuture = requestDataSupplier.sdkHttpResponseFuture();
-        return httpResponseFuture.thenCombine(putObjectOutputCompletableFuture,
-                                              (header, putObjectOutput) ->
-                                                  S3CrtPojoConversion.fromCrtPutObjectOutput(putObjectOutput, header));
+        CompletableFuture<PutObjectResponse> executeFuture =
+            httpResponseFuture.thenCombine(putObjectOutputCompletableFuture,
+                                           (header, putObjectOutput) -> S3CrtPojoConversion.fromCrtPutObjectOutput(
+                                               putObjectOutput, header));
+
+        return CompletableFutureUtils.forwardExceptionTo(executeFuture, putObjectOutputCompletableFuture);
     }
 
     @Override

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

@@ -23,12 +23,15 @@
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+import software.amazon.awssdk.transfer.s3.CompletedDownload;
+import software.amazon.awssdk.transfer.s3.CompletedUpload;
 import software.amazon.awssdk.transfer.s3.Download;
 import software.amazon.awssdk.transfer.s3.DownloadRequest;
 import software.amazon.awssdk.transfer.s3.S3ClientConfiguration;
 import software.amazon.awssdk.transfer.s3.S3TransferManager;
 import software.amazon.awssdk.transfer.s3.Upload;
 import software.amazon.awssdk.transfer.s3.UploadRequest;
+import software.amazon.awssdk.utils.CompletableFutureUtils;
 
 @SdkInternalApi
 public final class DefaultS3TransferManager implements S3TransferManager {
@@ -57,17 +60,21 @@ public Upload upload(UploadRequest uploadRequest) {
 
         CompletableFuture<PutObjectResponse> putObjFuture = s3CrtAsyncClient.putObject(putObjectRequest, requestBody);
 
-        return new DefaultUpload(putObjFuture.thenApply(r -> DefaultCompletedUpload.builder()
-                                                                                   .response(r)
-                                                                                   .build()));
+        CompletableFuture<CompletedUpload> future = putObjFuture.thenApply(r -> DefaultCompletedUpload.builder()
+                                                                                                      .response(r)
+                                                                                                      .build());
+        return new DefaultUpload(CompletableFutureUtils.forwardExceptionTo(future, putObjFuture));
     }
 
     @Override
     public Download download(DownloadRequest downloadRequest) {
-        CompletableFuture<GetObjectResponse> future =
+        CompletableFuture<GetObjectResponse> getObjectFuture =
             s3CrtAsyncClient.getObject(downloadRequest.getObjectRequest(),
                                        AsyncResponseTransformer.toFile(downloadRequest.destination()));
-        return new DefaultDownload(future.thenApply(r -> DefaultCompletedDownload.builder().response(r).build()));
+        CompletableFuture<CompletedDownload> future =
+            getObjectFuture.thenApply(r -> DefaultCompletedDownload.builder().response(r).build());
+
+        return new DefaultDownload(CompletableFutureUtils.forwardExceptionTo(future, getObjectFuture));
     }
 
     @Override

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

@@ -29,7 +29,7 @@
  * Internal client configuration resolver
  */
 @SdkInternalApi
-public final class S3NativeClientConfiguration implements SdkAutoCloseable {
+public class S3NativeClientConfiguration implements SdkAutoCloseable {
     private static final long DEFAULT_PART_SIZE_IN_BYTES = 8L * SizeConstant.MB;
     private static final long DEFAULT_TARGET_THROUGHPUT_IN_GBPS = 5;
     private final String signingRegion;

services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/DefaultS3CrtAsyncClientTest.java

@@ -0,0 +1,94 @@
+/*
+ * 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.Matchers.any;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.amazonaws.s3.RequestDataSupplier;
+import com.amazonaws.s3.ResponseDataConsumer;
+import com.amazonaws.s3.S3NativeClient;
+import com.amazonaws.s3.model.GetObjectOutput;
+import com.amazonaws.s3.model.GetObjectRequest;
+import com.amazonaws.s3.model.PutObjectOutput;
+import com.amazonaws.s3.model.PutObjectRequest;
+import java.util.concurrent.CompletableFuture;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.runners.MockitoJUnitRunner;
+import software.amazon.awssdk.core.ResponseBytes;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
+import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.services.s3.model.GetObjectResponse;
+import software.amazon.awssdk.services.s3.model.PutObjectResponse;
+
+@RunWith(MockitoJUnitRunner.class)
+public class DefaultS3CrtAsyncClientTest {
+    @Mock
+    private S3NativeClient mockS3NativeClient;
+
+    @Mock
+    private S3NativeClientConfiguration mockConfiguration;
+
+    private S3CrtAsyncClient s3CrtAsyncClient;
+
+    @Before
+    public void methodSetup() {
+        s3CrtAsyncClient = new DefaultS3CrtAsyncClient(mockConfiguration,
+                                                       mockS3NativeClient);
+    }
+
+    @Test
+    public void getObject_cancels_shouldForwardCancellation() {
+        CompletableFuture<GetObjectOutput> crtFuture = new CompletableFuture<>();
+        when(mockS3NativeClient.getObject(any(GetObjectRequest.class),
+                                          any(ResponseDataConsumer.class)))
+            .thenReturn(crtFuture);
+
+        CompletableFuture<ResponseBytes<GetObjectResponse>> future =
+            s3CrtAsyncClient.getObject(b -> b.bucket("bucket").key("key"),
+            AsyncResponseTransformer.toBytes());
+
+        future.cancel(true);
+        assertThat(crtFuture).isCancelled();
+    }
+
+    @Test
+    public void putObject_cancels_shouldForwardCancellation() {
+        CompletableFuture<PutObjectOutput> crtFuture = new CompletableFuture<>();
+        when(mockS3NativeClient.putObject(any(PutObjectRequest.class),
+                                          any(RequestDataSupplier.class)))
+            .thenReturn(crtFuture);
+
+        CompletableFuture<PutObjectResponse> future =
+            s3CrtAsyncClient.putObject(b -> b.bucket("bucket").key("key"),
+                                       AsyncRequestBody.empty());
+
+        future.cancel(true);
+        assertThat(crtFuture).isCancelled();
+    }
+
+    @Test
+    public void closeS3Client_shouldCloseUnderlyingResources() {
+        s3CrtAsyncClient.close();
+        verify(mockS3NativeClient).close();
+        verify(mockConfiguration).close();
+    }
+}

services-custom/s3-transfer-manager/src/test/java/software/amazon/awssdk/transfer/s3/internal/S3TransferManagerTest.java

@@ -73,6 +73,23 @@ public void upload_returnsResponse() {
         assertThat(completedUpload.response()).isEqualTo(response);
     }
 
+    @Test
+    public void upload_cancel_shouldForwardCancellation() {
+        CompletableFuture<PutObjectResponse> s3CrtFuture = new CompletableFuture<>();
+        when(mockS3Crt.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class)))
+            .thenReturn(s3CrtFuture);
+
+        CompletableFuture<CompletedUpload> future = tm.upload(UploadRequest.builder()
+                                                                                    .putObjectRequest(r -> r.bucket("bucket")
+                                                                                         .key("key"))
+                                                                                    .source(Paths.get("."))
+                                                                                    .build())
+                                                               .completionFuture();
+
+        future.cancel(true);
+        assertThat(s3CrtFuture).isCancelled();
+    }
+
     @Test
     public void download_returnsResponse() {
         GetObjectResponse response = GetObjectResponse.builder().build();
@@ -89,4 +106,21 @@ public void download_returnsResponse() {
         assertThat(completedDownload.response()).isEqualTo(response);
     }
 
+    @Test
+    public void download_cancel_shouldForwardCancellation() {
+        CompletableFuture<GetObjectResponse> s3CrtFuture = new CompletableFuture<>();
+        when(mockS3Crt.getObject(any(GetObjectRequest.class), any(AsyncResponseTransformer.class)))
+            .thenReturn(s3CrtFuture);
+
+        CompletableFuture<CompletedDownload> future = tm.download(DownloadRequest.builder()
+                                                                                 .getObjectRequest(r -> r.bucket("bucket")
+                                                                                                         .key("key"))
+                                                                                 .destination(Paths.get("."))
+                                                                                 .build())
+                                                        .completionFuture();
+        future.cancel(true);
+        assertThat(s3CrtFuture).isCancelled();
+    }
+
+
 }