From c5c822f18b65aac579479c820ef13cf949aac23a Mon Sep 17 00:00:00 2001
From: Dengke Tang <dengket@amazon.com>
Date: Sat, 2 Dec 2023 00:35:08 +0000
Subject: [PATCH] fixes

---
 source/s3_meta_request.c    | 12 ------------
 source/s3_request.c         | 19 ++++++++++++++++++-
 tests/s3_data_plane_tests.c |  2 ++
 3 files changed, 20 insertions(+), 13 deletions(-)

diff --git a/source/s3_meta_request.c b/source/s3_meta_request.c
index e00db911a..5c25c19aa 100644
--- a/source/s3_meta_request.c
+++ b/source/s3_meta_request.c
@@ -1386,18 +1386,6 @@ static void s_s3_meta_request_send_request_finish(
     struct aws_s3_meta_request_vtable *vtable = meta_request->vtable;
     AWS_PRECONDITION(vtable);
 
-    if (request->send_data.metrics) {
-        /* Invoke the callback here. */
-        struct aws_s3_request_metrics *metric = request->send_data.metrics;
-        aws_high_res_clock_get_ticks((uint64_t *)&metric->time_metrics.end_timestamp_ns);
-        metric->time_metrics.total_duration_ns =
-            metric->time_metrics.end_timestamp_ns - metric->time_metrics.start_timestamp_ns;
-        if (meta_request->telemetry_callback) {
-            meta_request->telemetry_callback(meta_request, metric, meta_request->user_data);
-        }
-        request->send_data.metrics = aws_s3_request_metrics_release(metric);
-    }
-
     vtable->send_request_finish(connection, stream, error_code);
 }
 
diff --git a/source/s3_request.c b/source/s3_request.c
index 29084b4f0..1c7d628d6 100644
--- a/source/s3_request.c
+++ b/source/s3_request.c
@@ -51,6 +51,22 @@ void aws_s3_request_setup_send_data(struct aws_s3_request *request, struct aws_h
     AWS_PRECONDITION(request);
     AWS_PRECONDITION(message);
 
+    if (request->send_data.metrics) {
+        /**
+         * When the request gets retried, finish up the metrics from the previous attampt.
+         */
+        struct aws_s3_request_metrics *metric = request->send_data.metrics;
+        aws_high_res_clock_get_ticks((uint64_t *)&metric->time_metrics.end_timestamp_ns);
+        metric->time_metrics.total_duration_ns =
+            metric->time_metrics.end_timestamp_ns - metric->time_metrics.start_timestamp_ns;
+
+        struct aws_s3_meta_request *meta_request = request->meta_request;
+        if (meta_request && meta_request->telemetry_callback) {
+            meta_request->telemetry_callback(meta_request, metric, meta_request->user_data);
+        }
+        request->send_data.metrics = aws_s3_request_metrics_release(metric);
+    }
+
     aws_s3_request_clean_up_send_data(request);
 
     request->send_data.message = message;
@@ -80,10 +96,11 @@ void aws_s3_request_clean_up_send_data(struct aws_s3_request *request) {
     s_s3_request_clean_up_send_data_message(request);
 
     aws_signable_destroy(request->send_data.signable);
-    request->send_data.signable = NULL;
+
     /* The metrics should be collected and provided before reaching here */
     AWS_FATAL_ASSERT(request->send_data.metrics == NULL);
 
+    request->send_data.signable = NULL;
     aws_http_headers_release(request->send_data.response_headers);
     request->send_data.response_headers = NULL;
 
diff --git a/tests/s3_data_plane_tests.c b/tests/s3_data_plane_tests.c
index e879a9557..d2545c099 100644
--- a/tests/s3_data_plane_tests.c
+++ b/tests/s3_data_plane_tests.c
@@ -404,6 +404,8 @@ static int s_test_s3_request_create_destroy(struct aws_allocator *allocator, voi
 
     request->send_data.response_headers = aws_http_headers_new(allocator);
     ASSERT_TRUE(request->send_data.response_headers != NULL);
+    ASSERT_TRUE(request->send_data.metrics != NULL);
+    request->send_data.metrics = aws_s3_request_metrics_release(request->send_data.metrics);
 
     aws_s3_request_clean_up_send_data(request);