Add response content length

aws · Jul 16, 2024 · f6f2c25 · f6f2c25
1 parent 3eb5af9
commit f6f2c25
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 8 deletions.
diff --git a/...n/java/software/amazon/awssdk/core/internal/http/pipeline/stages/HandleResponseStage.java b/...n/java/software/amazon/awssdk/core/internal/http/pipeline/stages/HandleResponseStage.java
@@ -15,6 +15,8 @@
 
 package software.amazon.awssdk.core.internal.http.pipeline.stages;
 
+import static software.amazon.awssdk.core.internal.util.ProgressListenerUtils.wrapContentStreamProviderWithByteReadTrackingIfProgressListenerAttached;
+
 import java.time.Duration;
 import java.util.concurrent.atomic.AtomicLong;
 import software.amazon.awssdk.annotations.SdkInternalApi;
@@ -24,9 +26,8 @@
 import software.amazon.awssdk.core.internal.http.RequestExecutionContext;
 import software.amazon.awssdk.core.internal.http.pipeline.RequestPipeline;
 import software.amazon.awssdk.core.internal.metrics.BytesReadTrackingInputStream;
-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.metrics.CoreMetric;
 import software.amazon.awssdk.http.AbortableInputStream;
 import software.amazon.awssdk.http.SdkHttpFullResponse;
@@ -47,6 +48,11 @@ public HandleResponseStage(HttpResponseHandler<Response<OutputT>> responseHandle
 
     @Override
     public Response<OutputT> execute(SdkHttpFullResponse httpResponse, RequestExecutionContext context) throws Exception {
+
+        context.progressUpdater().ifPresent(progressUpdater -> {
+            ProgressListenerUtils.updateProgressListenersWithResponseStatus(progressUpdater, httpResponse);
+        });
+
         SdkHttpFullResponse bytesReadTracking = trackBytesRead(httpResponse, context);
 
         Response<OutputT> response = responseHandler.handle(bytesReadTracking, context.executionAttributes());
@@ -88,12 +94,9 @@ private SdkHttpFullResponse trackBytesRead(SdkHttpFullResponse httpFullResponse,
     private AbortableInputStream trackBytesRead(AbortableInputStream content, RequestExecutionContext context) {
         AtomicLong bytesRead = context.executionAttributes().getAttribute(SdkInternalExecutionAttribute.RESPONSE_BYTES_READ);
 
-        ProgressUpdater progressUpdater = context.progressUpdater().isPresent() ?
-                                          context.progressUpdater().get() : null;
-
         BytesReadTrackingInputStream bytesReadTrackedStream =
-            new BytesReadTrackingInputStream(content, bytesRead,
-                                             new DownloadProgressUpdaterInvocation(progressUpdater));
+            wrapContentStreamProviderWithByteReadTrackingIfProgressListenerAttached(content, bytesRead, context);
+
         return AbortableInputStream.create(bytesReadTrackedStream);
     }
 }
diff --git a/.../java/software/amazon/awssdk/core/internal/progress/listener/LoggingProgressListener.java b/.../java/software/amazon/awssdk/core/internal/progress/listener/LoggingProgressListener.java
@@ -64,6 +64,7 @@ public void requestPrepared(Context.RequestPrepared context) {
 
     @Override
     public void requestHeaderSent(Context.RequestHeaderSent context) {
+        log.info(() -> "Started Upload...");
         context.uploadProgressSnapshot().ratioTransferred().ifPresent(progressBar::update);
     }
 

diff --git a/...k-core/src/main/java/software/amazon/awssdk/core/internal/util/ProgressListenerUtils.java b/...k-core/src/main/java/software/amazon/awssdk/core/internal/util/ProgressListenerUtils.java
@@ -57,7 +57,7 @@ public static ContentStreamProvider wrapContentStreamProviderWithBytePublishTrac
         if (progressUpdater != null) {
             wrappedContentStreamProvider =
                 ContentStreamProvider.fromInputStream(new BytesReadTrackingInputStream(
-                    (AbortableInputStream) contentStreamProvider.newStream(),
+                    AbortableInputStream.create(contentStreamProvider.newStream()),
                     new AtomicLong(0L),
                     new UploadProgressUpdaterInvocation(progressUpdater)));
         }

diff --git a/...a/software/amazon/awssdk/core/internal/progress/listener/LoggingProgressListenerTest.java b/...a/software/amazon/awssdk/core/internal/progress/listener/LoggingProgressListenerTest.java
@@ -65,6 +65,7 @@ public void defaultListener_successfulTransfer() {
             List<LogEvent> events = logCaptor.loggedEvents();
             assertLogged(events, Level.INFO, "Request Prepared...");
             assertLogged(events, Level.INFO, "|                    | 0.0%");
+            assertLogged(events, Level.INFO, "Started Upload...");
             assertLogged(events, Level.INFO, "|=                   | 5.0%");
             assertLogged(events, Level.INFO, "|==                  | 10.0%");
             assertLogged(events, Level.INFO, "|===                 | 15.0%");
@@ -120,6 +121,7 @@ public void test_customTicksListener_successfulTransfer() {
             List<LogEvent> events = logCaptor.loggedEvents();
             assertLogged(events, Level.INFO, "Request Prepared...");
             assertLogged(events, Level.INFO, "|     | 0.0%");
+            assertLogged(events, Level.INFO, "Started Upload...");
             assertLogged(events, Level.INFO, "|=    | 20.0%");
             assertLogged(events, Level.INFO, "|==   | 40.0%");
             assertLogged(events, Level.INFO, "|===  | 60.0%");