Fix PR comments, added an interface and abstract class to denote Upload and Download operations in a wya in which BytesReadTrackingPublisher can be reused for upload and download

Author: anirudh9391

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/internal/http/AmazonAsyncHttpClient.java

@@ -210,7 +210,7 @@ public <OutputT> CompletableFuture<OutputT> execute(
                                         .wrappedWith(AsyncApiCallAttemptMetricCollectionStage::new)
                                         .wrappedWith((deps, wrapped) -> new AsyncRetryableStage<>(responseHandler, deps, wrapped))
                                         .then(async(() -> new UnwrapResponseContainer<>()))
-                                        .then(() -> new AfterExecutionProgressReportingStage<>())
+                                        .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/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

@@ -25,7 +25,10 @@
 public class AfterExecutionProgressReportingStage<OutputT> implements RequestPipeline<OutputT, OutputT> {
     @Override
     public OutputT execute(OutputT input, RequestExecutionContext context) throws Exception {
-        ProgressListenerUtils.updateProgressListenersWithSuccessResponse((SdkResponse) input, context);
+        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/MakeAsyncHttpRequestStage.java

@@ -47,8 +47,11 @@
 import software.amazon.awssdk.core.internal.http.timers.TimeoutTracker;
 import software.amazon.awssdk.core.internal.http.timers.TimerUtils;
 import software.amazon.awssdk.core.internal.metrics.BytesReadTrackingPublisher;
+import software.amazon.awssdk.core.internal.progress.listener.ProgressUpdater;
+import software.amazon.awssdk.core.internal.util.DownloadProgressUpdaterInvocation;
 import software.amazon.awssdk.core.internal.util.MetricUtils;
 import software.amazon.awssdk.core.internal.util.ProgressListenerUtils;
+import software.amazon.awssdk.core.internal.util.UploadProgressUpdaterInvocation;
 import software.amazon.awssdk.core.metrics.CoreMetric;
 import software.amazon.awssdk.http.SdkHttpFullRequest;
 import software.amazon.awssdk.http.SdkHttpMethod;
@@ -130,22 +133,16 @@ private CompletableFuture<Response<OutputT>> executeHttpRequest(SdkHttpFullReque
 
         CompletableFuture<Response<OutputT>> responseHandlerFuture = responseHandler.prepare();
 
+        ProgressUpdater progressUpdater = ProgressListenerUtils.getProgressUpdaterIfAttached(context);
+
         SdkHttpContentPublisher requestProvider = context.requestProvider() == null
-                                                  ? new SimpleHttpContentPublisher(request)
-                                                  : new SdkHttpContentPublisherAdapter(context.requestProvider());
+                                                  ? new SimpleHttpContentPublisher(request, progressUpdater)
+                                                  : new SdkHttpContentPublisherAdapter(context.requestProvider(), progressUpdater);
         // Set content length if it hasn't been set already.
         SdkHttpFullRequest requestWithContentLength = getRequestWithContentLength(request, requestProvider);
 
         MetricCollector httpMetricCollector = MetricUtils.createHttpMetricsCollector(context);
 
-        //If Progress Listening is enabled, wrap around BytesSentTrackingPublisher to track progress on bytes sent
-        if (context.progressUpdater().isPresent()) {
-            boolean shouldSetContentLength = shouldSetContentLength(request, requestProvider);
-            requestProvider = ProgressListenerUtils.updateProgressListenersWithRequestStatus(context.progressUpdater().get(),
-                                                                                             requestProvider,
-                                                                                             shouldSetContentLength);
-        }
-
         AsyncExecuteRequest.Builder executeRequestBuilder = AsyncExecuteRequest.builder()
                                                                 .request(requestWithContentLength)
                                                                 .requestContentPublisher(requestProvider)
@@ -273,9 +270,11 @@ private void completeResponseFuture(CompletableFuture<Response<OutputT>> respons
     private static final class SdkHttpContentPublisherAdapter implements SdkHttpContentPublisher {
 
         private final AsyncRequestBody asyncRequestBody;
+        private final ProgressUpdater progressUpdater;
 
-        private SdkHttpContentPublisherAdapter(AsyncRequestBody asyncRequestBody) {
+        private SdkHttpContentPublisherAdapter(AsyncRequestBody asyncRequestBody, ProgressUpdater progressUpdater) {
             this.asyncRequestBody = asyncRequestBody;
+            this.progressUpdater = progressUpdater;
         }
 
         @Override
@@ -285,7 +284,12 @@ public Optional<Long> contentLength() {
 
         @Override
         public void subscribe(Subscriber<? super ByteBuffer> s) {
-            asyncRequestBody.subscribe(s);
+            if (progressUpdater != null) {
+                Publisher<ByteBuffer> readTrackingPublisher = new BytesReadTrackingPublisher(asyncRequestBody, new AtomicLong(0L), new UploadProgressUpdaterInvocation(progressUpdater));
+                readTrackingPublisher.subscribe(s);
+            } else {
+                asyncRequestBody.subscribe(s);
+            }
         }
     }
 
@@ -323,9 +327,11 @@ public void onStream(Publisher<ByteBuffer> stream) {
             AtomicLong bytesReadCounter = context.executionAttributes()
                                                  .getAttribute(SdkInternalExecutionAttribute.RESPONSE_BYTES_READ);
 
-            BytesReadTrackingPublisher bytesReadTrackingPublisher = new BytesReadTrackingPublisher(stream,
-                                                                        bytesReadCounter,
-                                                                        context.progressUpdater());
+            ProgressUpdater progressUpdater = context.progressUpdater().isPresent() ?
+                                              context.progressUpdater().get() : null;
+            Publisher<ByteBuffer> bytesReadTrackingPublisher = new BytesReadTrackingPublisher(stream,
+                                                                                              bytesReadCounter,
+                                                                                              new DownloadProgressUpdaterInvocation(progressUpdater));
 
             super.onStream(bytesReadTrackingPublisher);
         }

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

@@ -55,6 +55,7 @@ public Pair<SdkHttpFullRequest, SdkHttpFullResponse> execute(SdkHttpFullRequest
                                                                  RequestExecutionContext context) throws Exception {
         InterruptMonitor.checkInterrupted();
         HttpExecuteResponse executeResponse = executeHttpRequest(request, context);
+
         // TODO: Plumb through ExecuteResponse instead
         SdkHttpFullResponse httpResponse = (SdkHttpFullResponse) executeResponse.httpResponse();
         return Pair.of(request, httpResponse.toBuilder().content(executeResponse.responseBody().orElse(null)).build());

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

@@ -17,7 +17,6 @@
 
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.Response;
-import software.amazon.awssdk.core.SdkResponse;
 import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
 

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/metrics/BytesReadTrackingPublisher.java

@@ -16,13 +16,12 @@
 package software.amazon.awssdk.core.internal.metrics;
 
 import java.nio.ByteBuffer;
-import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
 import org.reactivestreams.Publisher;
 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.core.internal.util.ProgressUpdaterInvoker;
 
 /**
  * Publisher that tracks how many bytes are published from the wrapped publisher to the downstream subscriber.
@@ -31,20 +30,18 @@
 public final class BytesReadTrackingPublisher implements Publisher<ByteBuffer> {
     private final Publisher<ByteBuffer> upstream;
     private final AtomicLong bytesRead;
-    private ProgressUpdater progressUpdater;
+    private final ProgressUpdaterInvoker progressUpdaterInvoker;
 
     public BytesReadTrackingPublisher(Publisher<ByteBuffer> upstream, AtomicLong bytesRead,
-                                      Optional<ProgressUpdater> progressUpdater) {
+                                      ProgressUpdaterInvoker progressUpdaterInvoker) {
         this.upstream = upstream;
         this.bytesRead = bytesRead;
-        progressUpdater.ifPresent(value -> {
-            this.progressUpdater = value;
-        });
+        this.progressUpdaterInvoker = progressUpdaterInvoker;
     }
 
     @Override
     public void subscribe(Subscriber<? super ByteBuffer> subscriber) {
-        upstream.subscribe(new BytesReadTracker(subscriber, bytesRead, progressUpdater));
+        upstream.subscribe(new BytesReadTracker(subscriber, bytesRead, progressUpdaterInvoker));
     }
 
     public long bytesRead() {
@@ -54,29 +51,30 @@ public long bytesRead() {
     private static final class BytesReadTracker implements Subscriber<ByteBuffer> {
         private final Subscriber<? super ByteBuffer> downstream;
         private final AtomicLong bytesRead;
-        private final ProgressUpdater progressUpdater;
+        private final ProgressUpdaterInvoker progressUpdaterInvoker;
 
         private BytesReadTracker(Subscriber<? super ByteBuffer> downstream,
-                                 AtomicLong bytesRead, ProgressUpdater progressUpdater) {
+                                 AtomicLong bytesRead, ProgressUpdaterInvoker progressUpdaterInvoker) {
             this.downstream = downstream;
             this.bytesRead = bytesRead;
-            this.progressUpdater = progressUpdater;
+            this.progressUpdaterInvoker = progressUpdaterInvoker;
         }
 
         @Override
         public void onSubscribe(Subscription subscription) {
             downstream.onSubscribe(subscription);
-            if (progressUpdater != null) {
-                progressUpdater.resetBytesReceived();
+            if (progressUpdaterInvoker != null) {
+                progressUpdaterInvoker.resetBytes();
             }
         }
 
         @Override
         public void onNext(ByteBuffer byteBuffer) {
-            bytesRead.addAndGet(byteBuffer.remaining());
+            long byteBufferSize = byteBuffer.remaining();
+            bytesRead.addAndGet(byteBufferSize);
             downstream.onNext(byteBuffer);
-            if (progressUpdater != null) {
-                progressUpdater.incrementBytesReceived(bytesRead.get());
+            if (progressUpdaterInvoker != null) {
+                progressUpdaterInvoker.updateBytesTransferred(byteBufferSize);
             }
         }
 

core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/metrics/BytesSentTrackingPublisher.java

@@ -15,6 +15,7 @@
 
 package software.amazon.awssdk.core.internal.progress;
 
+import java.util.Optional;
 import software.amazon.awssdk.annotations.Immutable;
 import software.amazon.awssdk.annotations.SdkInternalApi;
 import software.amazon.awssdk.core.SdkRequest;
@@ -89,8 +90,8 @@ public SdkHttpResponse httpResponse() {
     }
 
     @Override
-    public SdkResponse response() {
-        return response;
+    public Optional<SdkResponse> response() {
+        return Optional.of(response);
     }
 
     @Override