Merge pull request #2679 from newrelic/add_sqs_instrumentation

Add SQS instrumentation

CHANGELOG.md

@@ -2,7 +2,11 @@
 
 ## dev
 
-Version <dev> fixes a bug related to expected errors not bearing a "true" value for the "expected" attribute if expected as a result of an HTTP status code match and changes the way Stripe instrumentation metrics are named to prevent high-cardinality issues.
+Version <dev> introduces instrumentation for the aws-sdk-sqs gem, fixes a bug related to expected errors not bearing a "true" value for the "expected" attribute if expected as a result of an HTTP status code match and changes the way Stripe instrumentation metrics are named to prevent high-cardinality issues.
+
+- **Feature: Add instrumentation for SQS**
+
+    The agent has added instrumentation for the [aws-sdk-sqs gem](https://rubygems.org/gems/aws-sdk-sqs). The agent will now record message broker spans for SQS client calls made with the aws-sdk-sqs gem. [PR#2679](https://github.com/newrelic/newrelic-ruby-agent/pull/2679)
 
 - **Bugfix: HTTP status code based expected errors will now have an "expected" value of "true"**
 

lib/new_relic/agent/configuration/default_source.rb

@@ -1458,6 +1458,14 @@ def self.enforce_fallback(allowed_values: nil, fallback: nil)
           :allowed_from_server => false,
           :description => 'Controls auto-instrumentation of bunny at start-up. May be one of: `auto`, `prepend`, `chain`, `disabled`.'
         },
+        :'instrumentation.aws_sqs' => {
+          :default => 'auto',
+          :public => true,
+          :type => String,
+          :dynamic_name => true,
+          :allowed_from_server => false,
+          :description => 'Controls auto-instrumentation of the aws-sdk-sqs library at start-up. May be one of: `auto`, `prepend`, `chain`, `disabled`.'
+        },
         :'instrumentation.dynamodb' => {
           :default => 'auto',
           :public => true,

lib/new_relic/agent/instrumentation/aws_sqs.rb

@@ -0,0 +1,25 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+require_relative 'aws_sqs/instrumentation'
+require_relative 'aws_sqs/chain'
+require_relative 'aws_sqs/prepend'
+
+DependencyDetection.defer do
+  named :aws_sqs
+
+  depends_on do
+    defined?(Aws::SQS::Client)
+  end
+
+  executes do
+    NewRelic::Agent.logger.info('Installing aws-sdk-sqs instrumentation')
+
+    if use_prepend?
+      prepend_instrument Aws::SQS::Client, NewRelic::Agent::Instrumentation::AwsSqs::Prepend
+    else
+      chain_instrument NewRelic::Agent::Instrumentation::AwsSqs::Chain
+    end
+  end
+end

lib/new_relic/agent/instrumentation/aws_sqs/chain.rb

@@ -0,0 +1,37 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+module NewRelic::Agent::Instrumentation
+  module AwsSqs::Chain
+    def self.instrument!
+      ::Aws::SQS::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::AwsSqs
+
+        alias_method(:send_message_without_new_relic, :send_message)
+
+        def send_message(*args)
+          send_message_with_new_relic(*args) do
+            send_message_without_new_relic(*args)
+          end
+        end
+
+        alias_method(:send_message_batch_without_new_relic, :send_message_batch)
+
+        def send_message_batch(*args)
+          send_message_batch_with_new_relic(*args) do
+            send_message_batch_without_new_relic(*args)
+          end
+        end
+
+        alias_method(:receive_message_without_new_relic, :receive_message)
+
+        def receive_message(*args)
+          receive_message_with_new_relic(*args) do
+            receive_message_without_new_relic(*args)
+          end
+        end
+      end
+    end
+  end
+end

lib/new_relic/agent/instrumentation/aws_sqs/instrumentation.rb

@@ -0,0 +1,67 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+module NewRelic::Agent::Instrumentation
+  module AwsSqs
+    MESSAGING_LIBRARY = 'SQS'
+
+    def send_message_with_new_relic(*args)
+      with_tracing(:produce, args) do
+        yield
+      end
+    end
+
+    def send_message_batch_with_new_relic(*args)
+      with_tracing(:produce, args) do
+        yield
+      end
+    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(params[0])
+        segment = NewRelic::Agent::Tracer.start_message_broker_segment(
+          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', e)
+      end
+      NewRelic::Agent::Tracer.capture_segment_error(segment) do
+        yield
+      end
+    ensure
+      segment&.finish
+    end
+
+    private
+
+    def add_aws_attributes(segment, info)
+      return unless segment
+
+      segment.add_agent_attribute('messaging.system', 'aws_sqs')
+      segment.add_agent_attribute('cloud.region', config&.region)
+      segment.add_agent_attribute('cloud.account.id', info[:account_id])
+      segment.add_agent_attribute('messaging.destination.name', info[:queue_name])
+    end
+
+    def get_url_info(params)
+      split = params[:queue_url].split('/')
+      {
+        queue_name: split.last,
+        account_id: split[-2]
+      }
+    end
+  end
+end

lib/new_relic/agent/instrumentation/aws_sqs/prepend.rb

@@ -0,0 +1,21 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+module NewRelic::Agent::Instrumentation
+  module AwsSqs::Prepend
+    include NewRelic::Agent::Instrumentation::AwsSqs
+
+    def send_message(*args)
+      send_message_with_new_relic(*args) { super }
+    end
+
+    def send_message_batch(*args)
+      send_message_batch_with_new_relic(*args) { super }
+    end
+
+    def receive_message(*args)
+      receive_message_with_new_relic(*args) { super }
+    end
+  end
+end

test/multiverse/suites/awssqs/Envfile

@@ -0,0 +1,10 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+instrumentation_methods :chain, :prepend
+
+gemfile <<~RB
+  gem 'aws-sdk-sqs'
+  gem 'nokogiri'
+RB

test/multiverse/suites/awssqs/awssqs_instrumentation_test.rb

@@ -0,0 +1,107 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+class AwssqsInstrumentationTest < Minitest::Test
+  def setup
+    Aws.config.update(stub_responses: true)
+  end
+
+  def teardown
+    harvest_span_events!
+    mocha_teardown
+  end
+
+  def create_client
+    Aws::SQS::Client.new(region: 'us-east-2')
+  end
+
+  def test_all_attributes_added_to_segment_send_message
+    client = create_client
+
+    in_transaction do |txn|
+      client.send_message({
+        queue_url: 'https://sqs.us-east-2.amazonaws.com/123456789/itsatestqueuewow',
+        message_body: 'wow, its a message'
+      })
+    end
+
+    spans = harvest_span_events!
+    span = spans[1][0]
+
+    assert_equal 'MessageBroker/SQS/Queue/Produce/Named/itsatestqueuewow', span[0]['name']
+
+    assert_equal 'aws_sqs', span[2]['messaging.system']
+    assert_equal 'us-east-2', span[2]['cloud.region']
+    assert_equal '123456789', span[2]['cloud.account.id']
+    assert_equal 'itsatestqueuewow', span[2]['messaging.destination.name']
+  end
+
+  def test_all_attributes_added_to_segment_send_message_batch
+    client = create_client
+
+    in_transaction do |txn|
+      client.send_message_batch({
+        queue_url: 'https://sqs.us-east-2.amazonaws.com/123456789/itsatestqueuewow',
+        entries: [
+          {
+            id: 'msq1',
+            message_body: 'wow 1'
+          },
+          {
+            id: 'msq2',
+            message_body: 'wow 2'
+          }
+        ]
+      })
+    end
+
+    spans = harvest_span_events!
+    span = spans[1][0]
+
+    assert_equal 'MessageBroker/SQS/Queue/Produce/Named/itsatestqueuewow', span[0]['name']
+
+    assert_equal 'aws_sqs', span[2]['messaging.system']
+    assert_equal 'us-east-2', span[2]['cloud.region']
+    assert_equal '123456789', span[2]['cloud.account.id']
+    assert_equal 'itsatestqueuewow', span[2]['messaging.destination.name']
+  end
+
+  def test_all_attributes_added_to_segment_receive_message
+    client = create_client
+
+    in_transaction do |txn|
+      client.receive_message({
+        queue_url: 'https://sqs.us-east-2.amazonaws.com/123456789/itsatestqueuewow'
+      })
+    end
+
+    spans = harvest_span_events!
+    span = spans[1][0]
+
+    assert_equal 'MessageBroker/SQS/Queue/Consume/Named/itsatestqueuewow', span[0]['name']
+
+    assert_equal 'aws_sqs', span[2]['messaging.system']
+    assert_equal 'us-east-2', span[2]['cloud.region']
+    assert_equal '123456789', span[2]['cloud.account.id']
+    assert_equal 'itsatestqueuewow', span[2]['messaging.destination.name']
+  end
+
+  def test_error_send_message
+    client = create_client
+
+    log = with_array_logger(:info) do
+      in_transaction do |txn|
+        begin
+          client.send_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')
+  end
+end

test/multiverse/suites/awssqs/config/newrelic.yml

@@ -0,0 +1,19 @@
+---
+development:
+  error_collector:
+    enabled: true
+  apdex_t: 0.5
+  monitor_mode: true
+  license_key: bootstrap_newrelic_admin_license_key_000
+  instrumentation:
+    aws_sqs: <%= $instrumentation_method %>
+  app_name: test
+  log_level: debug
+  host: 127.0.0.1
+  api_host: 127.0.0.1
+  transaction_trace:
+    record_sql: obfuscated
+    enabled: true
+    stack_trace_threshold: 0.5
+    transaction_threshold: 1.0
+  capture_params: false