From 73d7ce8a5f6509c5ad25ffc8851c57deb94764d4 Mon Sep 17 00:00:00 2001
From: Juli Tera <terajul@amazon.com>
Date: Thu, 22 Aug 2024 22:38:28 -0700
Subject: [PATCH] Move span attributes to http request and http response to
 make send middleware transport agnostic

---
 hearth/lib/hearth/http/request.rb    | 17 +++++++++++++
 hearth/lib/hearth/http/response.rb   | 13 ++++++++++
 hearth/lib/hearth/middleware/send.rb | 37 +++-------------------------
 3 files changed, 34 insertions(+), 33 deletions(-)

diff --git a/hearth/lib/hearth/http/request.rb b/hearth/lib/hearth/http/request.rb
index d267233d9..63705e8d1 100755
--- a/hearth/lib/hearth/http/request.rb
+++ b/hearth/lib/hearth/http/request.rb
@@ -154,6 +154,23 @@ def remove_query_param(name)
       def prefix_host(prefix)
         uri.host = prefix + uri.host
       end
+
+      # Contains attributes for Telemetry span to emit.
+      # @return [Hash]
+      def span_attributes
+        {
+          'http.method' => http_method,
+          'net.protocol.name' => 'http',
+          'net.protocol.version' => Net::HTTP::HTTPVersion,
+          'net.peer.name' => uri.host,
+          'net.peer.port' => uri.port
+        }.tap do |h|
+          if headers.key?('Content-Length')
+            h['http.request_content_length'] =
+              headers['Content-Length']
+          end
+        end
+      end
     end
   end
 end
diff --git a/hearth/lib/hearth/http/response.rb b/hearth/lib/hearth/http/response.rb
index eb26551bf..a72b95399 100755
--- a/hearth/lib/hearth/http/response.rb
+++ b/hearth/lib/hearth/http/response.rb
@@ -53,6 +53,19 @@ def reset
         @fields.clear
         super
       end
+
+      # Contains attributes for Telemetry span to emit.
+      # @return [Hash]
+      def span_attributes
+        {
+          'http.status_code' => status
+        }.tap do |h|
+          if headers.key?('Content-Length')
+            h['http.response_content_length'] =
+              headers['Content-Length']
+          end
+        end
+      end
     end
   end
 end
diff --git a/hearth/lib/hearth/middleware/send.rb b/hearth/lib/hearth/middleware/send.rb
index 5dfd198b0..58503b1be 100755
--- a/hearth/lib/hearth/middleware/send.rb
+++ b/hearth/lib/hearth/middleware/send.rb
@@ -80,7 +80,7 @@ def call(input, context)
       private
 
       def stub_response(input, context, output)
-        span_wrapper(context, stub_response: true) do
+        span_wrapper(context) do
           stub = @stubs.next(context.operation_name)
           log_debug(context, "Stubbing response with stub: #{stub}")
           if @response_events
@@ -109,42 +109,13 @@ def send_request(context, output)
         end
       end
 
-      def span_wrapper(context, stub_response: false, &block)
+      def span_wrapper(context, &block)
         context.tracer.in_span(
           'Middleware.Send',
-          attributes: request_attrs(context, stub_response: stub_response)
+          attributes: context.request.span_attributes
         ) do |span|
           block.call
-          span.add_attributes(response_attrs(context))
-        end
-      end
-
-      def request_attrs(context, stub_response: false)
-        {
-          'http.method' => context.request.http_method,
-          'net.protocol.name' => 'http',
-          'net.protocol.version' => Net::HTTP::HTTPVersion
-        }.tap do |h|
-          unless stub_response
-            h['net.peer.name'] = context.request.uri.host
-            h['net.peer.port'] = context.request.uri.port
-          end
-
-          if context.request.headers.key?('Content-Length')
-            h['http.request_content_length'] =
-              context.request.headers['Content-Length']
-          end
-        end
-      end
-
-      def response_attrs(context)
-        {
-          'http.status_code' => context.response.status
-        }.tap do |h|
-          if context.response.headers.key?('Content-Length')
-            h['http.response_content_length'] =
-              context.response.headers['Content-Length']
-          end
+          span.add_attributes(context.response.span_attributes)
         end
       end