Update HTTP Async SPI

This commit makes changes to the HTTP Async SPI as well as the
AsyncResponseTransformer interface exposed by SDK clients.

For the HTTP SPI, the main changes are to SdkHttpResponseHandler and
SdkAsyncHttpClient.

- SdkHttpResponseHandler

  The most important change is to complete(). It has been removed as it was a
  potential spot where the async client could get blocked or where the
  implementation would be forced to block, since the handler must return the
  transformed response from this method. We also don't necessarily care about
  what the response handler does at the HTTP client layer, so there's no need to
  return the transformed result.

  Secondly, the interface now guarantees that onStream() is always called, even
  for responses with no content. This should simplify implementation a bit since
  implementers know to always expect a stream.

  Other changes are mostly minor, such as naming: the interface is now
  SdkAsyncResponseHandler to be more intuitive since it follows our current
  naming conventions.

- SdkAsyncHttpClient

  prepareRequest() has been replaced with execute() that takes a single
  parameter object. It now also returns a CompletableFuture rather than a
  RunnableRequest.

- AsyncResponseTransformer

  This continues to be almost identical to SdkAsyncHttpResponseHandler.

  The most important change here is in complete(). Like in
  SdkHttpResponseHandler, this method was problematic since it had the
  protential to block the caller. To address this, complete() has been replaced
  with a method named prepare() that returns a CompletableFuture. We guarantee
  that this method is always called before executing the request to allow
  implementers to do any setup/cleanup. This method could be called more than
  once in the event of retries.

  Another important result of prepare() returning a CompletableFuture is this
  gives users the ability to trigger retries by completing it exceptionally with
  a instance of RetryableException similar to the sync ReponseTransformer.

  Finally, this change decouples response transformation from request
  completion. With the current behavior, the complete() is only called when the
  entire response is received. This makes doing things like returning the
  Publisher from onStream() as the result of the transformer awkward.

Author: dagnir

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/AsyncClientClass.java

@@ -179,7 +179,7 @@ protected MethodSpec.Builder operationBody(MethodSpec.Builder builder, Operation
         if (opModel.hasStreamingOutput() || opModel.hasEventStreamOutput()) {
             String paramName = opModel.hasStreamingOutput() ? "asyncResponseTransformer" : "asyncResponseHandler";
             builder.addStatement("runAndLogError(log, \"Exception thrown in exceptionOccurred callback, ignoring\",\n" +
-                                 "() -> $L.exceptionOccurred(t))", paramName);
+                                 "() -> $N.exceptionOccurred(t))", paramName);
         }
 
         return builder.addStatement("return $T.failedFuture(t)", CompletableFutureUtils.class)

codegen/src/main/java/software/amazon/awssdk/codegen/poet/client/specs/QueryXmlProtocolSpec.java

@@ -154,20 +154,14 @@ public CodeBlock asyncExecutionHandler(IntermediateModel intermediateModel, Oper
                                        ".withResponseHandler(responseHandler)" +
                                        ".withErrorResponseHandler($N)\n" +
                                        asyncRequestBody +
-                                       ".withInput($L) $L)$L;",
+                                       ".withInput($L) $L);",
                                        ClientExecutionParams.class,
                                        requestType,
                                        pojoResponseType,
                                        marshaller,
                                        "errorResponseHandler",
                                        opModel.getInput().getVariableName(),
-                                       opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "",
-                                       // If it's a streaming operation we also need to notify the handler on exception
-                                       opModel.hasStreamingOutput() ? ".whenComplete((r, e) -> {\n"
-                                                                      + "    if (e != null) {\n"
-                                                                      + "        asyncResponseTransformer.exceptionOccurred(e);\n"
-                                                                      + "    }\n"
-                                                                      + "})" : "")
+                                       opModel.hasStreamingOutput() ? ", asyncResponseTransformer" : "")
                         .build();
     }
 

codegen/src/test/resources/software/amazon/awssdk/codegen/poet/client/test-async-client-class.java

@@ -282,8 +282,7 @@ public CompletableFuture<Void> eventStreamOperation(EventStreamOperationRequest
             });
             return future;
         } catch (Throwable t) {
-            runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring",
-                           () -> asyncResponseHandler.exceptionOccurred(t));
+            runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring", () -> asyncResponseHandler.exceptionOccurred(t));
             return CompletableFutureUtils.failedFuture(t);
         }
     }
@@ -764,7 +763,7 @@ public <ReturnT> CompletableFuture<ReturnT> streamingOutputOperation(
             });
         } catch (Throwable t) {
             runAndLogError(log, "Exception thrown in exceptionOccurred callback, ignoring",
-                           () -> asyncResponseTransformer.exceptionOccurred(t));
+                    () -> asyncResponseTransformer.exceptionOccurred(t));
             return CompletableFutureUtils.failedFuture(t);
         }
     }

core/aws-core/src/main/java/software/amazon/awssdk/awscore/eventstream/EventStreamAsyncResponseTransformer.java

@@ -46,8 +46,8 @@
 import software.amazon.awssdk.core.http.HttpResponseHandler;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
-import software.amazon.awssdk.core.internal.util.ThrowableUtils;
 import software.amazon.awssdk.http.AbortableInputStream;
+import software.amazon.awssdk.http.SdkCancellationException;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
 import software.amazon.awssdk.utils.BinaryUtils;
 import software.amazon.eventstream.Message;
@@ -114,11 +114,6 @@ public class EventStreamAsyncResponseTransformer<ResponseT, EventT>
      */
     private volatile boolean isDone = false;
 
-    /**
-     * Holds a reference to any exception delivered to exceptionOccurred.
-     */
-    private final AtomicReference<Throwable> error = new AtomicReference<>();
-
     /**
      * Executor to deliver events to the subscriber
      */
@@ -161,6 +156,8 @@ public class EventStreamAsyncResponseTransformer<ResponseT, EventT>
      */
     private String requestId = null;
 
+    private volatile CompletableFuture<Void> transformFuture;
+
     @Deprecated
     @ReviewBeforeRelease("Remove this on full GA of 2.0.0")
     public EventStreamAsyncResponseTransformer(
@@ -191,7 +188,16 @@ private EventStreamAsyncResponseTransformer(
     }
 
     @Override
-    public void responseReceived(SdkResponse response) {
+    public CompletableFuture<Void> prepare() {
+        transformFuture = new CompletableFuture<>();
+        return transformFuture;
+    }
+
+    @Override
+    public void onResponse(SdkResponse response) {
+        // We use a void unmarshaller and unmarshall the actual response in the message
+        // decoder when we receive the initial-response frame. TODO not clear
+        // how we would handle REST protocol which would unmarshall the response from the HTTP headers
         if (response != null && response.sdkHttpResponse() != null) {
             this.requestId = response.sdkHttpResponse()
                                      .firstMatchingHeader(X_AMZN_REQUEST_ID_HEADER)
@@ -223,31 +229,17 @@ public void exceptionOccurred(Throwable throwable) {
         synchronized (this) {
             if (!isDone) {
                 isDone = true;
-                error.set(throwable);
                 // If we have a Subscriber at this point notify it as well
-                if (subscriberRef.get() != null) {
+                if (subscriberRef.get() != null && shouldSurfaceErrorToEventSubscriber(throwable)) {
                     runAndLogError(log, "Error thrown from Subscriber#onError, ignoring.",
                         () -> subscriberRef.get().onError(throwable));
                 }
                 eventStreamResponseHandler.exceptionOccurred(throwable);
+                transformFuture.completeExceptionally(throwable);
             }
         }
     }
 
-    @Override
-    public Void complete() {
-        if (error.get() == null) {
-            // Add the special on complete event to signal drainEvents to complete the subscriber
-            eventsToDeliver.add(ON_COMPLETE_EVENT);
-            drainEventsIfNotAlready();
-            return null;
-        } else {
-            // Need to propagate the failure up so the future is completed exceptionally. This should only happen
-            // when there is a frame level exception that the upper layers don't know about.
-            throw ThrowableUtils.failure(error.get());
-        }
-    }
-
     /**
      * Called when all events have been delivered to the downstream subscriber.
      */
@@ -340,6 +332,10 @@ private SdkHttpFullResponse adaptMessageToResponse(Message message, boolean isEx
                                   .build();
     }
 
+    private static boolean shouldSurfaceErrorToEventSubscriber(Throwable t) {
+        return !(t instanceof SdkCancellationException);
+    }
+
     /**
      * Subscriber for the raw bytes from the stream. Feeds them to the {@link MessageDecoder} as they arrive
      * and will request as much as needed to fulfill any outstanding demand.
@@ -390,8 +386,10 @@ public void onError(Throwable throwable) {
 
         @Override
         public void onComplete() {
-            // Notified in onEventComplete method because we have more context on what we've delivered to
-            // the event stream subscriber there.
+            // Add the special on complete event to signal drainEvents to complete the subscriber
+            eventsToDeliver.add(ON_COMPLETE_EVENT);
+            drainEventsIfNotAlready();
+            transformFuture.complete(null);
         }
     }
 

core/aws-core/src/main/java/software/amazon/awssdk/awscore/eventstream/RestEventStreamAsyncResponseTransformer.java

@@ -16,6 +16,8 @@
 package software.amazon.awssdk.awscore.eventstream;
 
 import java.nio.ByteBuffer;
+import java.util.concurrent.CompletableFuture;
+
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
@@ -43,8 +45,13 @@ private RestEventStreamAsyncResponseTransformer(
     }
 
     @Override
-    public void responseReceived(ResponseT response) {
-        delegate.responseReceived(response);
+    public CompletableFuture<Void> prepare() {
+        return delegate.prepare();
+    }
+
+    @Override
+    public void onResponse(ResponseT response) {
+        delegate.onResponse(response);
         eventStreamResponseHandler.responseReceived(response);
     }
 
@@ -58,11 +65,6 @@ public void exceptionOccurred(Throwable throwable) {
         delegate.exceptionOccurred(throwable);
     }
 
-    @Override
-    public Void complete() {
-        return delegate.complete();
-    }
-
     public static <ResponseT extends SdkResponse, EventT> Builder<ResponseT, EventT> builder() {
         return new Builder<>();
     }

core/aws-core/src/test/java/software/amazon/awssdk/awscore/eventstream/EventStreamAsyncResponseTransformerTest.java

@@ -23,6 +23,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -123,10 +124,19 @@ private void verifyExceptionThrown(Map<String, HeaderValue> headers) {
                                                .executor(Executors.newSingleThreadExecutor())
                                                .future(new CompletableFuture<>())
                                                .build();
-        transformer.responseReceived(null);
+        CompletableFuture<Void> cf = transformer.prepare();
+        transformer.onResponse(null);
         transformer.onStream(SdkPublisher.adapt(bytePublisher));
 
-        assertThatThrownBy(transformer::complete).isSameAs(exception);
+        assertThatThrownBy(() -> {
+            try {
+                cf.join();
+            } catch (CompletionException e) {
+                if (e.getCause() instanceof SdkServiceException) {
+                    throw ((SdkServiceException) e.getCause());
+                }
+            }
+        }).isSameAs(exception);
     }
 
     private static class SubscribingResponseHandler implements EventStreamResponseHandler<Object, Object> {

core/sdk-core/src/main/java/software/amazon/awssdk/core/async/AsyncResponseTransformer.java

@@ -18,70 +18,92 @@
 import java.io.File;
 import java.nio.ByteBuffer;
 import java.nio.file.Path;
-import org.reactivestreams.Publisher;
-import org.reactivestreams.Subscription;
+import java.util.concurrent.CompletableFuture;
+
 import software.amazon.awssdk.annotations.SdkPublicApi;
 import software.amazon.awssdk.core.ResponseBytes;
 import software.amazon.awssdk.core.internal.async.ByteArrayAsyncResponseTransformer;
 import software.amazon.awssdk.core.internal.async.FileAsyncResponseTransformer;
 
 /**
  * Callback interface to handle a streaming asynchronous response.
+ * <p>
+ * <h2>Synchronization</h2>
+ * <p>
+ * All operations, including those called on the {@link org.reactivestreams.Subscriber} of the stream are guaranteed to be
+ * synchronized externally; i.e. no two methods on this interface or on the {@link org.reactivestreams.Subscriber} will be
+ * invoked concurrently. It is <b>not</b> guaranteed that the methods will being invoked by the same thread.
+ * <p>
+ * <h2>Invocation Order</h2>
+ * <p>
+ * The methods are called in the following order:
+ * <ul>
+ *     <li>
+ *     {@link #prepare()}: This method is always called first. Implementations should use this to setup or perform any
+ *     cleanup necessary. <b>Note that this will be called upon each request attempt</b>. If the {@link CompletableFuture}
+ *     returned from the previous invocation has already been completed, the implementation should return a new instance.
+ *     </li>
+ *     <li>
+ *     {@link #onResponse}: If the response was received successfully, this method is called next.
+ *     </li>
+ *     <li>
+ *     {@link #onStream(SdkPublisher)}: Called after {@code onResponse}. This is always invoked, even if the service
+ *     operation response does not contain a body. If the response does not have a body, then the {@link SdkPublisher} will
+ *     complete the subscription without signaling any elements.
+ *     </li>
+ *     <li>
+ *     {@link #exceptionOccurred(Throwable)}: If there is an error sending the request. This method is called before {@link
+ *     org.reactivestreams.Subscriber#onError(Throwable)}.
+ *     </li>
+ *     <li>
+ *     {@link org.reactivestreams.Subscriber#onError(Throwable)}: If an error is encountered while the {@code Publisher} is
+ *     publishing to a {@link org.reactivestreams.Subscriber}.
+ *     </li>
+ * </ul>
+ * <p>
+ * <h2>Retries</h2>
+ * <p>
+ * The transformer has the ability to trigger retries at any time by completing the {@link CompletableFuture} with an
+ * exception that is deemed retryable by the configured {@link software.amazon.awssdk.core.retry.RetryPolicy}.
  *
  * @param <ResponseT> POJO response type.
- * @param <ReturnT>   Type this response handler produces. I.E. the type you are transforming the response into.
+ * @param <ResultT>   Type this response handler produces. I.E. the type you are transforming the response into.
  */
 @SdkPublicApi
-public interface AsyncResponseTransformer<ResponseT, ReturnT> {
-
+public interface AsyncResponseTransformer<ResponseT, ResultT> {
     /**
-     * Called when the initial response has been received and the POJO response has
-     * been unmarshalled. This is guaranteed to be called before onStream.
-     *
-     * <p>In the event of a retryable error, this callback may be called multiple times. It
-     * also may never be invoked if the request never succeeds.</p>
+     * Initial call to enable any setup required before the response is handled.
+     * <p>
+     * Note that this will be called for each request attempt, up to the number of retries allowed by the configured {@link
+     * software.amazon.awssdk.core.retry.RetryPolicy}.
+     * <p>
+     * This method is guaranteed to be called before the request is executed, and before {@link #onResponse(Object)} is
+     * signaled.
      *
-     * @param response Unmarshalled POJO containing metadata about the streamed data.
+     * @return The future holding the transformed response.
      */
-    void responseReceived(ResponseT response);
+    CompletableFuture<ResultT> prepare();
 
     /**
-     * Called when events are ready to be streamed. Implementations  must subscribe to the {@link Publisher} and request data via
-     * a {@link org.reactivestreams.Subscription} as they can handle it.
+     * Called when the unmarshalled response object is ready.
      *
-     * <p>
-     * If at any time the subscriber wishes to stop receiving data, it may call {@link Subscription#cancel()}. This
-     * will be treated as a failure of the response and the {@link #exceptionOccurred(Throwable)} callback will be invoked.
-     * </p>
-     *
-     * <p>This callback may never be called if the response has no content or if an error occurs.</p>
-     *
-     * <p>
-     * In the event of a retryable error, this callback may be called multiple times with different Publishers.
-     * If this method is called more than once, implementation must either reset any state to prepare for another
-     * stream of data or must throw an exception indicating they cannot reset. If any exception is thrown then no
-     * automatic retry is performed.
-     * </p>
+     * @param response The unmarshalled response.
      */
-    void onStream(SdkPublisher<ByteBuffer> publisher);
+    void onResponse(ResponseT response);
 
     /**
-     * Called when an exception occurs while establishing the connection or streaming the response. Implementations
-     * should free up any resources in this method. This method may be called multiple times during the lifecycle
-     * of a request if automatic retries are enabled.
+     * Called when the response stream is ready.
      *
-     * @param throwable Exception that occurred.
+     * @param publisher The publisher.
      */
-    void exceptionOccurred(Throwable throwable);
+    void onStream(SdkPublisher<ByteBuffer> publisher);
 
     /**
-     * Called when all data has been successfully published to the {@link org.reactivestreams.Subscriber}. This will
-     * only be called once during the lifecycle of the request. Implementors should free up any resources they have
-     * opened and do final transformations to produce the return object.
+     * Called when a error is encountered while making the request or receiving the response.
      *
-     * @return Transformed object as a result of the streamed data.
+     * @param error Error that occurred.
      */
-    ReturnT complete();
+    void exceptionOccurred(Throwable error);
 
     /**
      * Creates an {@link AsyncResponseTransformer} that writes all the content to the given file. In the event of an error,