Invoke responseHandler#onError when the request future is cancelled

Author: zoewangg

http-clients/aws-crt-client/src/main/java/software/amazon/awssdk/http/crt/internal/CrtRequestExecutor.java

@@ -30,6 +30,7 @@
 import software.amazon.awssdk.crt.http.HttpHeader;
 import software.amazon.awssdk.crt.http.HttpRequest;
 import software.amazon.awssdk.http.Header;
+import software.amazon.awssdk.http.SdkCancellationException;
 import software.amazon.awssdk.http.SdkHttpRequest;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpResponseHandler;
@@ -41,7 +42,7 @@ public final class CrtRequestExecutor {
     private static final Logger log = Logger.loggerFor(CrtRequestExecutor.class);
 
     public CompletableFuture<Void> execute(CrtRequestContext executionContext) {
-        CompletableFuture<Void> requestFuture = new CompletableFuture<>();
+        CompletableFuture<Void> requestFuture = createExecutionFuture(executionContext.sdkRequest());
 
         // When a Connection is ready from the Connection Pool, schedule the Request on the connection
         CompletableFuture<HttpClientConnection> httpClientConnectionCompletableFuture =
@@ -77,6 +78,27 @@ public CompletableFuture<Void> execute(CrtRequestContext executionContext) {
         return requestFuture;
     }
 
+    /**
+     * Convenience method to create the execution future and set up the cancellation logic.
+     *
+     * @return The created execution future.
+     */
+    private CompletableFuture<Void> createExecutionFuture(AsyncExecuteRequest request) {
+        CompletableFuture<Void> future = new CompletableFuture<>();
+
+        future.whenComplete((r, t) -> {
+            if (t == null) {
+                return;
+            }
+            //TODO: Aborting request once it's supported in CRT
+            if (future.isCancelled()) {
+                request.responseHandler().onError(new SdkCancellationException("The request was cancelled"));
+            }
+        });
+
+        return future;
+    }
+
     private void handleFailure(Throwable cause,
                                CompletableFuture<Void> executeFuture,
                                SdkAsyncHttpResponseHandler responseHandler) {

http-clients/aws-crt-client/src/test/java/software/amazon/awssdk/http/crt/internal/CrtRequestExecutorTest.java

@@ -22,6 +22,7 @@
 import java.io.IOException;
 import java.net.URI;
 import java.util.concurrent.CompletableFuture;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -33,6 +34,7 @@
 import software.amazon.awssdk.crt.http.HttpClientConnection;
 import software.amazon.awssdk.crt.http.HttpClientConnectionManager;
 import software.amazon.awssdk.crt.http.HttpRequest;
+import software.amazon.awssdk.http.SdkCancellationException;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpResponseHandler;
@@ -43,6 +45,7 @@ public class CrtRequestExecutorTest {
     private CrtRequestExecutor requestExecutor;
     @Mock
     private HttpClientConnectionManager connectionManager;
+
     @Mock
     private SdkAsyncHttpResponseHandler responseHandler;
 
@@ -54,8 +57,13 @@ public void setup() {
         requestExecutor = new CrtRequestExecutor();
     }
 
+    @After
+    public void teardown() {
+        Mockito.reset(connectionManager, responseHandler, httpClientConnection);
+    }
+
     @Test
-    public void execute_acquireConnectionThrowException_shouldInvokeOnError() {
+    public void acquireConnectionThrowException_shouldInvokeOnError() {
         RuntimeException exception = new RuntimeException("error");
         CrtRequestContext context = CrtRequestContext.builder()
                                                      .crtConnPool(connectionManager)
@@ -80,7 +88,7 @@ public void execute_acquireConnectionThrowException_shouldInvokeOnError() {
     }
 
     @Test
-    public void execute_makeRequestThrowException_shouldInvokeOnError() {
+    public void makeRequestThrowException_shouldInvokeOnError() {
         CrtRuntimeException exception = new CrtRuntimeException("");
         SdkHttpFullRequest request = createRequest(URI.create("http://localhost"));
         CrtRequestContext context = CrtRequestContext.builder()
@@ -110,4 +118,47 @@ public void execute_makeRequestThrowException_shouldInvokeOnError() {
         assertThat(actualException).hasCause(exception);
         assertThat(executeFuture).hasFailedWithThrowableThat().hasCause(exception).isInstanceOf(IOException.class);
     }
+
+    @Test
+    public void makeRequest_success() {
+        SdkHttpFullRequest request = createRequest(URI.create("http://localhost"));
+        CrtRequestContext context = CrtRequestContext.builder()
+                                                     .readBufferSize(2000)
+                                                     .crtConnPool(connectionManager)
+                                                     .request(AsyncExecuteRequest.builder()
+                                                                                 .request(request)
+                                                                                 .requestContentPublisher(createProvider(""))
+                                                                                 .responseHandler(responseHandler)
+                                                                                 .build())
+                                                     .build();
+        CompletableFuture<HttpClientConnection> completableFuture = new CompletableFuture<>();
+        Mockito.when(connectionManager.acquireConnection()).thenReturn(completableFuture);
+        completableFuture.complete(httpClientConnection);
+
+        CompletableFuture<Void> executeFuture = requestExecutor.execute(context);
+        Mockito.verifyZeroInteractions(responseHandler);
+    }
+
+    @Test
+    public void cancelRequest_shouldInvokeOnError() {
+        CrtRequestContext context = CrtRequestContext.builder()
+                                                     .crtConnPool(connectionManager)
+                                                     .request(AsyncExecuteRequest.builder()
+                                                                                 .responseHandler(responseHandler)
+                                                                                 .build())
+                                                     .build();
+        CompletableFuture<HttpClientConnection> completableFuture = new CompletableFuture<>();
+
+        Mockito.when(connectionManager.acquireConnection()).thenReturn(completableFuture);
+
+        CompletableFuture<Void> executeFuture = requestExecutor.execute(context);
+        executeFuture.cancel(true);
+
+        ArgumentCaptor<Exception> argumentCaptor = ArgumentCaptor.forClass(Exception.class);
+        Mockito.verify(responseHandler).onError(argumentCaptor.capture());
+
+        Exception actualException = argumentCaptor.getValue();
+        assertThat(actualException).hasMessageContaining("The request was cancelled");
+        assertThat(actualException).isInstanceOf(SdkCancellationException.class);
+    }
 }

test/http-client-tests/pom.xml

@@ -48,6 +48,11 @@
             <artifactId>http-client-spi</artifactId>
             <version>${awsjavasdk.version}</version>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>metrics-spi</artifactId>
+            <version>${awsjavasdk.version}</version>
+        </dependency>
         <dependency>
             <groupId>software.amazon.awssdk</groupId>
             <artifactId>utils</artifactId>

test/http-client-tests/src/main/java/software/amazon/awssdk/http/RecordingResponseHandler.java

@@ -34,7 +34,6 @@ public final class RecordingResponseHandler implements SdkAsyncHttpResponseHandl
 
     @Override
     public void onHeaders(SdkHttpResponse response) {
-        System.out.println("received response");
         responses.add(response);
     }