refactor

newrelic · Jun 20, 2024 · df6d3ca · df6d3ca
1 parent 278a3b5
commit df6d3ca
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 72 deletions.
diff --git a/lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb b/lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb
@@ -7,60 +7,36 @@ module AwsSqs
     MESSAGING_LIBRARY = 'SQS'
 
     def send_message_with_new_relic(*args)
-      segment = nil
-      begin
-        info = get_url_info(args[0])
-        segment = NewRelic::Agent::Tracer.start_message_broker_segment(
-          action: :produce,
-          library: MESSAGING_LIBRARY,
-          destination_type: :queue,
-          destination_name: info[:queue_name]
-        )
-        add_aws_attributes(segment, info)
-      rescue => e
-        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client#send_message ', e)
-      end
-      NewRelic::Agent::Tracer.capture_segment_error(segment) do
+      with_tracing(:produce, args) do
         yield
       end
-    ensure
-      segment&.finish
     end
 
     def send_message_batch_with_new_relic(*args)
-      segment = nil
-      begin
-        info = get_url_info(args[0])
-        segment = NewRelic::Agent::Tracer.start_message_broker_segment(
-          action: :produce,
-          library: MESSAGING_LIBRARY,
-          destination_type: :queue,
-          destination_name: info[:queue_name]
-        )
-        add_aws_attributes(segment, info)
-      rescue => e
-        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client#send_message_batch ', e)
-      end
-      NewRelic::Agent::Tracer.capture_segment_error(segment) do
+      with_tracing(:produce, args) do
         yield
       end
-    ensure
-      segment&.finish
     end
 
     def receive_message_with_new_relic(*args)
+      with_tracing(:consume, args) do
+        yield
+      end
+    end
+
+    def with_tracing(action, params)
       segment = nil
       begin
-        info = get_url_info(args[0])
+        info = get_url_info(params[0])
         segment = NewRelic::Agent::Tracer.start_message_broker_segment(
-          action: :consume,
+          action: action,
           library: MESSAGING_LIBRARY,
           destination_type: :queue,
           destination_name: info[:queue_name]
         )
         add_aws_attributes(segment, info)
       rescue => e
-        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client#receive_message ', e)
+        NewRelic::Agent.logger.error('Error starting message broker segment in Aws::SQS::Client', e)
       end
       NewRelic::Agent::Tracer.capture_segment_error(segment) do
         yield
@@ -69,6 +45,8 @@ def receive_message_with_new_relic(*args)
       segment&.finish
     end
 
+    private
+
     def add_aws_attributes(segment, info)
       return unless segment
 

diff --git a/test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb b/test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb
@@ -102,42 +102,6 @@ def test_error_send_message
       end
     end
 
-    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client#send_message')
-  end
-
-  def test_error_send_message_batch
-    client = create_client
-
-    log = with_array_logger(:info) do
-      in_transaction do |txn|
-        begin
-          client.send_message_batch({
-            queue_url: 42
-          })
-        rescue
-          # will cause an error in the instrumentation, but also will make the sdk raise an error
-        end
-      end
-    end
-
-    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client#send_message_batch')
-  end
-
-  def test_error_receive_message
-    client = create_client
-
-    log = with_array_logger(:info) do
-      in_transaction do |txn|
-        begin
-          client.receive_message({
-            queue_url: 42
-          })
-        rescue
-          # will cause an error in the instrumentation, but also will make the sdk raise an error
-        end
-      end
-    end
-
-    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client#receive_message')
+    assert_log_contains(log, 'Error starting message broker segment in Aws::SQS::Client')
   end
 end