Fixed the issue where the SDK did not offload future completion to the future completion executor

Author: zoewangg

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/pipeline/stages/MakeAsyncHttpRequestStage.java

@@ -26,7 +26,6 @@
 import java.util.concurrent.Executor;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.function.Supplier;
-import org.reactivestreams.Publisher;
 import org.reactivestreams.Subscriber;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.Response;
@@ -48,7 +47,6 @@
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpMethod;
-import software.amazon.awssdk.http.SdkHttpResponse;
 import software.amazon.awssdk.http.async.AsyncExecuteRequest;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
@@ -119,59 +117,12 @@ public CompletableFuture<Response<OutputT>> execute(CompletableFuture<SdkHttpFul
         return toReturn;
     }
 
-    private static final class WrappedErrorForwardingResponseHandler<T>
-            implements TransformingAsyncResponseHandler<T> {
-
-        private final TransformingAsyncResponseHandler<T> wrappedHandler;
-        private final CompletableFuture<T> responseFuture;
-
-        private WrappedErrorForwardingResponseHandler(TransformingAsyncResponseHandler<T> wrappedHandler,
-                                                      CompletableFuture<T> responseFuture) {
-            this.wrappedHandler = wrappedHandler;
-            this.responseFuture = responseFuture;
-
-        }
-
-        private static <T> WrappedErrorForwardingResponseHandler<T> of(
-                TransformingAsyncResponseHandler<T> wrappedHandler,
-                CompletableFuture<T> responseFuture) {
-
-            return new WrappedErrorForwardingResponseHandler<>(wrappedHandler, responseFuture);
-        }
-
-        @Override
-        public CompletableFuture<T> prepare() {
-            return wrappedHandler.prepare();
-        }
-
-        @Override
-        public void onHeaders(SdkHttpResponse headers) {
-            wrappedHandler.onHeaders(headers);
-        }
-
-        @Override
-        public void onStream(Publisher<ByteBuffer> stream) {
-            wrappedHandler.onStream(stream);
-        }
-
-        @Override
-        public void onError(Throwable error) {
-            responseFuture.completeExceptionally(error);
-            wrappedHandler.onError(error);
-        }
-    }
-
     private CompletableFuture<Response<OutputT>> executeHttpRequest(SdkHttpFullRequest request,
                                                                     RequestExecutionContext context) {
 
         CompletableFuture<Response<OutputT>> responseFuture = new CompletableFuture<>();
 
-        // Wrap the response handler in a layer that will notify the newly created responseFuture when the onError event
-        // is triggered
-        TransformingAsyncResponseHandler<Response<OutputT>> wrappedResponseHandler =
-            WrappedErrorForwardingResponseHandler.of(responseHandler, responseFuture);
-
-        CompletableFuture<Response<OutputT>> responseHandlerFuture = wrappedResponseHandler.prepare();
+        CompletableFuture<Response<OutputT>> responseHandlerFuture = responseHandler.prepare();
 
         SdkHttpContentPublisher requestProvider = context.requestProvider() == null
                                                   ? new SimpleHttpContentPublisher(request)
@@ -184,7 +135,7 @@ private CompletableFuture<Response<OutputT>> executeHttpRequest(SdkHttpFullReque
         AsyncExecuteRequest.Builder executeRequestBuilder = AsyncExecuteRequest.builder()
                                                                 .request(requestWithContentLength)
                                                                 .requestContentPublisher(requestProvider)
-                                                                .responseHandler(wrappedResponseHandler)
+                                                                .responseHandler(responseHandler)
                                                                 .fullDuplex(isFullDuplex(context.executionAttributes()))
                                                                 .metricCollector(httpMetricCollector);
         if (context.executionAttributes().getAttribute(SDK_HTTP_EXECUTION_ATTRIBUTES) != null) {
@@ -224,10 +175,10 @@ private CompletableFuture<Void> doExecuteHttpRequest(RequestExecutionContext con
         CompletableFuture<Void> httpClientFuture = sdkAsyncHttpClient.execute(executeRequest);
 
         // Offload the metrics reporting from this stage onto the future completion executor
-        CompletableFuture<Void> result = httpClientFuture.whenComplete((r, t) -> {
+        CompletableFuture<Void> result = httpClientFuture.whenCompleteAsync((r, t) -> {
             long duration = System.nanoTime() - callStart;
             metricCollector.reportMetric(CoreMetric.SERVICE_CALL_DURATION, Duration.ofNanos(duration));
-        });
+        }, futureCompletionExecutor);
 
         // Make sure failures on the result future are forwarded to the http client future.
         CompletableFutureUtils.forwardExceptionTo(result, httpClientFuture);

test/protocol-tests/src/test/java/software/amazon/awssdk/protocol/tests/AsyncResponseThreadingTest.java

@@ -21,10 +21,13 @@
 import static com.github.tomakehurst.wiremock.client.WireMock.urlPathEqualTo;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThat;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.Mockito.atLeast;
+import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
+import static org.mockito.Matchers.any;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption.FUTURE_COMPLETION_EXECUTOR;
 
+import com.github.tomakehurst.wiremock.http.Fault;
 import com.github.tomakehurst.wiremock.junit.WireMockRule;
 import java.net.URI;
 import java.util.concurrent.Executor;
@@ -34,8 +37,11 @@
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
+import software.amazon.awssdk.core.exception.SdkClientException;
+import software.amazon.awssdk.core.retry.RetryPolicy;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.protocolrestjson.ProtocolRestJsonAsyncClient;
+import software.amazon.awssdk.services.protocolrestjson.model.ProtocolRestJsonException;
 import software.amazon.awssdk.services.protocolrestjson.model.StreamingOutputOperationRequest;
 import software.amazon.awssdk.services.protocolrestjson.model.StreamingOutputOperationResponse;
 
@@ -64,12 +70,58 @@ public void completionWithNioThreadWorksCorrectly() {
                                                 AsyncResponseTransformer.toBytes()).join();
 
         // #1 reporting metrics, #2 completing response
-        verify(mockExecutor, atLeast(1)).execute(any());
+        verify(mockExecutor, times(2)).execute(any());
 
         byte[] arrayCopy = response.asByteArray();
         assertThat(arrayCopy).containsExactly('t', 'e', 's', 't');
     }
 
+    @Test
+    public void connectionError_completionWithNioThreadWorksCorrectly() {
+        stubFor(post(urlPathEqualTo(STREAMING_OUTPUT_PATH)).willReturn(aResponse().withFault(Fault.CONNECTION_RESET_BY_PEER).withBody("test")));
+
+        Executor mockExecutor = Mockito.spy(new SpyableExecutor());
+
+        ProtocolRestJsonAsyncClient client =
+            ProtocolRestJsonAsyncClient.builder()
+                                       .region(Region.US_WEST_1)
+                                       .endpointOverride(URI.create("http://localhost:" + wireMock.port()))
+                                       .credentialsProvider(() -> AwsBasicCredentials.create("akid", "skid"))
+                                       .asyncConfiguration(c -> c.advancedOption(FUTURE_COMPLETION_EXECUTOR, mockExecutor))
+                                       .overrideConfiguration(o -> o.retryPolicy(RetryPolicy.none()))
+                                       .build();
+
+        assertThatThrownBy(() ->
+            client.streamingOutputOperation(StreamingOutputOperationRequest.builder().build(),
+                                            AsyncResponseTransformer.toBytes()).join())
+            .hasCauseInstanceOf(SdkClientException.class);
+
+        // #1 reporting metrics, #2 completing response
+        verify(mockExecutor, times(2)).execute(any());
+    }
+
+    @Test
+    public void serverError_completionWithNioThreadWorksCorrectly() {
+        stubFor(post(urlPathEqualTo(STREAMING_OUTPUT_PATH)).willReturn(aResponse().withStatus(500).withBody("test")));
+
+        Executor mockExecutor = Mockito.spy(new SpyableExecutor());
+
+        ProtocolRestJsonAsyncClient client =
+            ProtocolRestJsonAsyncClient.builder()
+                                       .region(Region.US_WEST_1)
+                                       .endpointOverride(URI.create("http://localhost:" + wireMock.port()))
+                                       .credentialsProvider(() -> AwsBasicCredentials.create("akid", "skid"))
+                                       .overrideConfiguration(o -> o.retryPolicy(RetryPolicy.none()))
+                                       .asyncConfiguration(c -> c.advancedOption(FUTURE_COMPLETION_EXECUTOR, mockExecutor))
+                                       .build();
+
+        assertThatThrownBy(() ->
+            client.streamingOutputOperation(StreamingOutputOperationRequest.builder().build(),
+                                            AsyncResponseTransformer.toBytes()).join()).hasCauseInstanceOf(ProtocolRestJsonException.class);
+        // #1 reporting metrics, #2 completing response
+        verify(mockExecutor, times(2)).execute(any());
+    }
+
     private static class SpyableExecutor implements Executor {
         @Override
         public void execute(Runnable command) {

test/protocol-tests/src/test/java/software/amazon/awssdk/protocol/tests/retry/AsyncAwsJsonRetryTest.java

@@ -51,7 +51,7 @@ public class AsyncAwsJsonRetryTest {
     public void setupClient() {
         client = ProtocolJsonRpcAsyncClient.builder()
                                            .credentialsProvider(StaticCredentialsProvider.create(AwsBasicCredentials.create
-                                               ("akid", "skid")))
+                                                                                                                        ("akid", "skid")))
                                            .region(Region.US_EAST_1)
                                            .endpointOverride(URI.create("http://localhost:" + wireMock.port()))
                                            .build();
@@ -155,4 +155,4 @@ public void retryPolicyNone_shouldNotRetry() {
         assertThatThrownBy(() -> clientWithNoRetry.allTypes(AllTypesRequest.builder().build()).join())
             .hasCauseInstanceOf(ProtocolJsonRpcException.class);
     }
-}
+}

test/protocol-tests/src/test/resources/jetty-logging.properties

@@ -15,4 +15,4 @@
 
 # Set up logging implementation
 org.eclipse.jetty.util.log.class=org.eclipse.jetty.util.log.StdErrLog
-org.eclipse.jetty.LEVEL=INFO
+org.eclipse.jetty.LEVEL=OFF