Add Progress Listener Invocation methods to Asynchronous and Synchronous code paths

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.AsyncRetryableStage2;
 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,13 +203,15 @@ 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 AsyncRetryableStage2<>(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

@@ -27,13 +27,15 @@
 import software.amazon.awssdk.core.http.HttpResponseHandler;
 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;
@@ -190,7 +192,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)
@@ -204,6 +207,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

@@ -26,6 +26,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 +45,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 +129,14 @@ public void attemptMetricCollector(MetricCollector metricCollector) {
         this.attemptMetricCollector = metricCollector;
     }
 
+    public ProgressUpdater progressUpdater() {
+        return progressUpdater;
+    }
+
+    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/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.progressUpdater());
+        }
+
+        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,17 @@ 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
+                reportFailureToProgressListeners(context.progressUpdater(), 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,47 @@
+/*
+ * 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.DefaultProgressUpdater;
+import software.amazon.awssdk.core.internal.progress.listener.NoOpProgressUpdater;
+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 DefaultProgressUpdater(context.originalRequest(), requestContentLength);
+            progressUpdater.requestPrepared(input);
+            context.progressUpdater(progressUpdater);
+        } else {
+            context.progressUpdater(new NoOpProgressUpdater());
+        }
+
+        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,8 @@ public OutputT execute(SdkHttpFullRequest input, RequestExecutionContext context
             return wrapped.execute(input, context);
         } catch (Exception e) {
             Throwable throwable = reportFailureToInterceptors(context, e);
+
+            reportFailureToProgressListeners(context.progressUpdater(), throwable);
             throw failure(throwable);
         }
     }

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

@@ -15,6 +15,9 @@
 
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
+import static software.amazon.awssdk.core.internal.util.ProgressListenerUtils.updateProgressListenersWithResponseStatus;
+import static software.amazon.awssdk.core.internal.util.ProgressListenerUtils.wrapWithBytesReadTrackingStream;
+
 import java.time.Duration;
 import java.util.concurrent.atomic.AtomicLong;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -45,10 +48,11 @@ public HandleResponseStage(HttpResponseHandler<Response<OutputT>> responseHandle
 
     @Override
     public Response<OutputT> execute(SdkHttpFullResponse httpResponse, RequestExecutionContext context) throws Exception {
+
+        updateProgressListenersWithResponseStatus(context.progressUpdater(), httpResponse);
         SdkHttpFullResponse bytesReadTracking = trackBytesRead(httpResponse, context);
 
         Response<OutputT> response = responseHandler.handle(bytesReadTracking, context.executionAttributes());
-
         collectMetrics(context);
 
         return response;
@@ -85,7 +89,11 @@ private SdkHttpFullResponse trackBytesRead(SdkHttpFullResponse httpFullResponse,
 
     private AbortableInputStream trackBytesRead(AbortableInputStream content, RequestExecutionContext context) {
         AtomicLong bytesRead = context.executionAttributes().getAttribute(SdkInternalExecutionAttribute.RESPONSE_BYTES_READ);
-        BytesReadTrackingInputStream bytesReadTrackedStream = new BytesReadTrackingInputStream(content, bytesRead);
+
+        BytesReadTrackingInputStream bytesReadTrackedStream =
+            wrapWithBytesReadTrackingStream(
+                content, bytesRead, context.progressUpdater());
+
         return AbortableInputStream.create(bytesReadTrackedStream);
     }
 }