Add Progress Listener Invocation methods to Asynchronous code paths

core/imds/src/main/java/software/amazon/awssdk/imds/internal/AsyncHttpRequestHelper.java

@@ -61,7 +61,7 @@ private static <T> CompletableFuture<T> sendAsync(SdkAsyncHttpClient client,
                                                       SdkHttpFullRequest request,
                                                       HttpResponseHandler<T> handler,
                                                       CompletableFuture<?> parentFuture) {
-        SdkHttpContentPublisher requestContentPublisher = new SimpleHttpContentPublisher(request);
+        SdkHttpContentPublisher requestContentPublisher = new SimpleHttpContentPublisher(request, null);
         TransformingAsyncResponseHandler<T> responseHandler =
             new AsyncResponseHandler<>(handler, Function.identity(), new ExecutionAttributes());
         CompletableFuture<T> responseHandlerFuture = responseHandler.prepare();

core/sdk-core/src/main/java/software/amazon/awssdk/core/http/ExecutionContext.java

@@ -15,13 +15,11 @@
 
 package software.amazon.awssdk.core.http;
 
-import java.util.Optional;
 import software.amazon.awssdk.annotations.NotThreadSafe;
 import software.amazon.awssdk.annotations.SdkProtectedApi;
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.interceptor.ExecutionInterceptorChain;
 import software.amazon.awssdk.core.interceptor.InterceptorContext;
-import software.amazon.awssdk.core.internal.progress.listener.ProgressUpdater;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.metrics.MetricCollector;
 import software.amazon.awssdk.utils.builder.CopyableBuilder;
@@ -38,15 +36,13 @@ public final class ExecutionContext implements ToCopyableBuilder<ExecutionContex
     private final ExecutionInterceptorChain interceptorChain;
     private final ExecutionAttributes executionAttributes;
     private final MetricCollector metricCollector;
-    private final ProgressUpdater progressUpdater;
 
     private ExecutionContext(final Builder builder) {
         this.signer = builder.signer;
         this.interceptorContext = builder.interceptorContext;
         this.interceptorChain = builder.interceptorChain;
         this.executionAttributes = builder.executionAttributes;
         this.metricCollector = builder.metricCollector;
-        this.progressUpdater = builder.progressUpdater;
     }
 
     public static ExecutionContext.Builder builder() {
@@ -80,10 +76,6 @@ public MetricCollector metricCollector() {
         return metricCollector;
     }
 
-    public Optional<ProgressUpdater> progressUpdater() {
-        return progressUpdater != null ? Optional.of(progressUpdater) : Optional.empty();
-    }
-
     @Override
     public Builder toBuilder() {
         return new Builder(this);
@@ -95,7 +87,6 @@ public static class Builder implements CopyableBuilder<Builder, ExecutionContext
         private ExecutionAttributes executionAttributes;
         private Signer signer;
         private MetricCollector metricCollector;
-        private ProgressUpdater progressUpdater;
 
         private Builder() {
         }
@@ -106,7 +97,6 @@ private Builder(ExecutionContext executionContext) {
             this.interceptorChain = executionContext.interceptorChain;
             this.executionAttributes = executionContext.executionAttributes;
             this.metricCollector = executionContext.metricCollector;
-            this.progressUpdater = executionContext.progressUpdater;
         }
 
         public Builder interceptorContext(InterceptorContext interceptorContext) {
@@ -134,11 +124,6 @@ public Builder metricCollector(MetricCollector metricCollector) {
             return this;
         }
 
-        public Builder progressUpdater(ProgressUpdater progressUpdater) {
-            this.progressUpdater = progressUpdater;
-            return this;
-        }
-
         @Override
         public ExecutionContext build() {
             return new ExecutionContext(this);

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonAsyncHttpClient.java

@@ -30,6 +30,7 @@
 import software.amazon.awssdk.core.http.ExecutionContext;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipelineBuilder;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AfterExecutionInterceptorsStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.AfterExecutionProgressReportingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApplyTransactionIdStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApplyUserAgentStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncApiCallAttemptMetricCollectionStage;
@@ -39,6 +40,7 @@
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncExecutionFailureExceptionReportingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncRetryableStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AsyncSigningStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.BeforeExecutionProgressReportingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.CompressRequestStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.HttpChecksumStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.MakeAsyncHttpRequestStage;
@@ -201,12 +203,14 @@ public <OutputT> CompletableFuture<OutputT> execute(
                                 .then(() -> new HttpChecksumStage(ClientType.ASYNC))
                                 .then(MakeRequestImmutableStage::new)
                                 .then(RequestPipelineBuilder
-                                        .first(AsyncSigningStage::new)
+                                        .first(BeforeExecutionProgressReportingStage::new)
+                                        .then(AsyncSigningStage::new)
                                         .then(AsyncBeforeTransmissionExecutionInterceptorsStage::new)
                                         .then(d -> new MakeAsyncHttpRequestStage<>(responseHandler, d))
                                         .wrappedWith(AsyncApiCallAttemptMetricCollectionStage::new)
                                         .wrappedWith((deps, wrapped) -> new AsyncRetryableStage<>(responseHandler, deps, wrapped))
                                         .then(async(() -> new UnwrapResponseContainer<>()))
+                                        .then(async(() -> new AfterExecutionProgressReportingStage<>()))
                                         .then(async(() -> new AfterExecutionInterceptorsStage<>()))
                                         .wrappedWith(AsyncExecutionFailureExceptionReportingStage::new)
                                         .wrappedWith(AsyncApiCallTimeoutTrackingStage::new)

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/AmazonSyncHttpClient.java

@@ -28,13 +28,15 @@
 import software.amazon.awssdk.core.interceptor.ExecutionAttributes;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipelineBuilder;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AfterExecutionInterceptorsStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.AfterExecutionProgressReportingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.AfterTransmissionExecutionInterceptorsStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApiCallAttemptMetricCollectionStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApiCallAttemptTimeoutTrackingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApiCallMetricCollectionStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApiCallTimeoutTrackingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApplyTransactionIdStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.ApplyUserAgentStage;
+import software.amazon.awssdk.core.internal.http.pipeline.stages.BeforeExecutionProgressReportingStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.BeforeTransmissionExecutionInterceptorsStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.BeforeUnmarshallingExecutionInterceptorsStage;
 import software.amazon.awssdk.core.internal.http.pipeline.stages.CompressRequestStage;
@@ -204,7 +206,8 @@ public <OutputT> OutputT execute(HttpResponseHandler<Response<OutputT>> response
                                .then(MakeRequestImmutableStage::new)
                                // End of mutating request
                                .then(RequestPipelineBuilder
-                                         .first(SigningStage::new)
+                                         .first(BeforeExecutionProgressReportingStage::new)
+                                         .then(SigningStage::new)
                                          .then(BeforeTransmissionExecutionInterceptorsStage::new)
                                          .then(MakeHttpRequestStage::new)
                                          .then(AfterTransmissionExecutionInterceptorsStage::new)
@@ -218,6 +221,7 @@ public <OutputT> OutputT execute(HttpResponseHandler<Response<OutputT>> response
                                .wrappedWith(ApiCallTimeoutTrackingStage::new)::build)
                                .wrappedWith((deps, wrapped) -> new ApiCallMetricCollectionStage<>(wrapped))
                     .then(() -> new UnwrapResponseContainer<>())
+                    .then(() -> new AfterExecutionProgressReportingStage<>())
                     .then(() -> new AfterExecutionInterceptorsStage<>())
                     .wrappedWith(ExecutionFailureExceptionReportingStage::new)
                     .build(httpClientDependencies)

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/RequestExecutionContext.java

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.core.internal.http;
 
+import java.util.Optional;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.RequestOverrideConfiguration;
 import software.amazon.awssdk.core.SdkRequest;
@@ -26,6 +27,7 @@
 import software.amazon.awssdk.core.interceptor.SdkExecutionAttribute;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
 import software.amazon.awssdk.core.internal.http.timers.TimeoutTracker;
+import software.amazon.awssdk.core.internal.progress.listener.ProgressUpdater;
 import software.amazon.awssdk.core.signer.Signer;
 import software.amazon.awssdk.metrics.MetricCollector;
 import software.amazon.awssdk.utils.Validate;
@@ -44,6 +46,7 @@ public final class RequestExecutionContext {
     private TimeoutTracker apiCallTimeoutTracker;
     private TimeoutTracker apiCallAttemptTimeoutTracker;
     private MetricCollector attemptMetricCollector;
+    private ProgressUpdater progressUpdater;
 
     private RequestExecutionContext(Builder builder) {
         this.requestProvider = builder.requestProvider;
@@ -127,6 +130,14 @@ public void attemptMetricCollector(MetricCollector metricCollector) {
         this.attemptMetricCollector = metricCollector;
     }
 
+    public Optional<ProgressUpdater> progressUpdater() {
+        return progressUpdater != null ? Optional.of(progressUpdater) : Optional.empty();
+    }
+
+    public void progressUpdater(ProgressUpdater progressUpdater) {
+        this.progressUpdater = progressUpdater;
+    }
+
     /**
      * Sets the request body provider.
      * Used for transforming the original body provider to sign events for

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/http/async/SimpleHttpContentPublisher.java

@@ -22,6 +22,7 @@
 import org.reactivestreams.Subscriber;
 import org.reactivestreams.Subscription;
 import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.internal.progress.listener.ProgressUpdater;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.async.SdkHttpContentPublisher;
 import software.amazon.awssdk.utils.IoUtils;
@@ -35,11 +36,13 @@ public final class SimpleHttpContentPublisher implements SdkHttpContentPublisher
 
     private final byte[] content;
     private final int length;
+    private final ProgressUpdater progressUpdater;
 
-    public SimpleHttpContentPublisher(SdkHttpFullRequest request) {
+    public SimpleHttpContentPublisher(SdkHttpFullRequest request, ProgressUpdater progressUpdater) {
         this.content = request.contentStreamProvider().map(p -> invokeSafely(() -> IoUtils.toByteArray(p.newStream())))
                                                       .orElseGet(() -> new byte[0]);
         this.length = content.length;
+        this.progressUpdater = progressUpdater;
     }
 
     @Override
@@ -52,7 +55,7 @@ public void subscribe(Subscriber<? super ByteBuffer> s) {
         s.onSubscribe(new SubscriptionImpl(s));
     }
 
-    private class SubscriptionImpl implements Subscription {
+    private final class SubscriptionImpl implements Subscription {
         private boolean running = true;
         private final Subscriber<? super ByteBuffer> s;
 
@@ -68,6 +71,11 @@ public void request(long n) {
                     s.onError(new IllegalArgumentException("Demand must be positive"));
                 } else {
                     s.onNext(ByteBuffer.wrap(content));
+
+                    if (progressUpdater != null) {
+                        progressUpdater.incrementBytesSent(length);
+                    }
+
                     s.onComplete();
                 }
             }

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

@@ -0,0 +1,36 @@
+/*
+ * 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.core.internal.http.pipeline.stages;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.SdkResponse;
+import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
+import software.amazon.awssdk.core.internal.util.ProgressListenerUtils;
+
+@SdkInternalApi
+public class AfterExecutionProgressReportingStage<OutputT> implements RequestPipeline<OutputT, OutputT> {
+    @Override
+    public OutputT execute(OutputT input, RequestExecutionContext context) throws Exception {
+        if (input instanceof SdkResponse) {
+            ProgressListenerUtils.updateProgressListenersWithSuccessResponse((SdkResponse) input, context);
+        }
+
+        return input;
+    }
+
+
+}

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

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
 import static software.amazon.awssdk.core.internal.http.pipeline.stages.utils.ExceptionReportingUtils.reportFailureToInterceptors;
+import static software.amazon.awssdk.core.internal.http.pipeline.stages.utils.ExceptionReportingUtils.reportFailureToProgressListeners;
 
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -42,16 +43,19 @@ public CompletableFuture<OutputT> execute(SdkHttpFullRequest input, RequestExecu
         CompletableFuture<OutputT> executeFuture = wrappedExecute.handle((o, t) -> {
             if (t != null) {
                 Throwable toReport = t;
-
                 if (toReport instanceof CompletionException) {
                     toReport = toReport.getCause();
                 }
                 toReport = reportFailureToInterceptors(context, toReport);
 
+                // If Progress Listeners are attached to the request, update them with the throwable
+                if (context.progressUpdater().isPresent()) {
+                    reportFailureToProgressListeners(context.progressUpdater().get(), toReport);
+                }
+
                 throw CompletableFutureUtils.errorAsCompletionException(ThrowableUtils.asSdkException(toReport));
-            } else {
-                return o;
             }
+            return o;
         });
         return CompletableFutureUtils.forwardExceptionTo(executeFuture, wrappedExecute);
     }

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

@@ -0,0 +1,43 @@
+/*
+ * 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.core.internal.http.pipeline.stages;
+
+import software.amazon.awssdk.annotations.SdkInternalApi;
+import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
+import software.amazon.awssdk.core.internal.http.pipeline.RequestToRequestPipeline;
+import software.amazon.awssdk.core.internal.progress.listener.ProgressUpdater;
+import software.amazon.awssdk.core.internal.util.ProgressListenerUtils;
+import software.amazon.awssdk.http.SdkHttpFullRequest;
+
+@SdkInternalApi
+public class BeforeExecutionProgressReportingStage implements RequestToRequestPipeline {
+
+    @Override
+    public SdkHttpFullRequest execute(SdkHttpFullRequest input, RequestExecutionContext context) throws Exception {
+
+        if (ProgressListenerUtils.progressListenerAttached(context.originalRequest())) {
+            Long requestContentLength =
+                context.requestProvider() != null && context.requestProvider().contentLength().isPresent() ?
+                context.requestProvider().contentLength().get() : null;
+
+            ProgressUpdater progressUpdater = new ProgressUpdater(context.originalRequest(), requestContentLength);
+            progressUpdater.requestPrepared(input);
+            context.progressUpdater(progressUpdater);
+        }
+
+        return input;
+    }
+}

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

@@ -16,6 +16,7 @@
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
 import static software.amazon.awssdk.core.internal.http.pipeline.stages.utils.ExceptionReportingUtils.reportFailureToInterceptors;
+import static software.amazon.awssdk.core.internal.http.pipeline.stages.utils.ExceptionReportingUtils.reportFailureToProgressListeners;
 import static software.amazon.awssdk.core.internal.util.ThrowableUtils.failure;
 
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -37,6 +38,10 @@ public OutputT execute(SdkHttpFullRequest input, RequestExecutionContext context
             return wrapped.execute(input, context);
         } catch (Exception e) {
             Throwable throwable = reportFailureToInterceptors(context, e);
+
+            context.progressUpdater().ifPresent(progressUpdater -> {
+                reportFailureToProgressListeners(progressUpdater, throwable);
+            });
             throw failure(throwable);
         }
     }