Dynamodb instrumentation

lib/new_relic/agent.rb

@@ -64,6 +64,7 @@ module Agent
     require 'new_relic/agent/linking_metadata'
     require 'new_relic/agent/local_log_decorator'
     require 'new_relic/agent/llm'
+    require 'new_relic/agent/aws'
 
     require 'new_relic/agent/instrumentation/controller_instrumentation'
 

lib/new_relic/agent/aws.rb

@@ -0,0 +1,56 @@
+# 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
+  module Agent
+    module Aws
+      CHARACTERS = %w[A B C D E F G H I J K L M N O P Q R S T U V W X Y Z 2 3 4 5 6 7].freeze
+      HEX_MASK = '7fffffffff80'
+
+      def self.create_arn(service, resource, config)
+        region = config.region
+        account_id = NewRelic::Agent::Aws.convert_access_key_to_account_id(config.credentials.access_key_id)
+
+        "arn:aws:#{service}:#{region}:#{account_id}:#{resource}"
+      rescue => e
+        NewRelic::Agent.logger.warn("Failed to create ARN: #{e}")
+      end
+
+      def self.convert_access_key_to_account_id(access_key)
+        decoded_key = Integer(decode_to_hex(access_key[4..-1]), 16)
+        mask = Integer(HEX_MASK, 16)
+        (decoded_key & mask) >> 7
+      end
+
+      def self.decode_to_hex(access_key)
+        bytes = access_key.delete('=').each_char.map { |c| CHARACTERS.index(c) }
+
+        bytes.each_slice(8).map do |section|
+          convert_section(section)
+        end.flatten[0...6].join
+      end
+
+      def self.convert_section(section)
+        buffer = 0
+        section.each do |chunk|
+          buffer = (buffer << 5) + chunk
+        end
+
+        chunk_count = (section.length * 5.0 / 8.0).floor
+
+        if section.length < 8
+          buffer >>= (5 - (chunk_count * 8)) % 5
+        end
+
+        decoded = []
+        chunk_count.times do |i|
+          shift = 8 * (chunk_count - 1 - i)
+          decoded << ((buffer >> shift) & 255).to_s(16)
+        end
+
+        decoded
+      end
+    end
+  end
+end

lib/new_relic/agent/configuration/default_source.rb

@@ -1451,6 +1451,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.dynamodb' => {
+          :default => 'auto',
+          :public => true,
+          :type => String,
+          :dynamic_name => true,
+          :allowed_from_server => false,
+          :description => 'Controls auto-instrumentation of the aws-sdk-dynamodb library at start-up. May be one of [auto|prepend|chain|disabled].'
+        },
         :'instrumentation.fiber' => {
           :default => 'auto',
           :public => true,

lib/new_relic/agent/instrumentation/dynamodb.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 'dynamodb/instrumentation'
+require_relative 'dynamodb/chain'
+require_relative 'dynamodb/prepend'
+
+DependencyDetection.defer do
+  named :dynamodb
+
+  depends_on do
+    defined?(Aws::DynamoDB::Client)
+  end
+
+  executes do
+    NewRelic::Agent.logger.info('Installing dynamodb instrumentation')
+
+    if use_prepend?
+      prepend_instrument Aws::DynamoDB::Client, NewRelic::Agent::Instrumentation::Dynamodb::Prepend
+    else
+      chain_instrument NewRelic::Agent::Instrumentation::Dynamodb::Chain
+    end
+  end
+end

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

@@ -0,0 +1,34 @@
+# 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 Dynamodb::Chain
+    def self.instrument!
+      ::Aws::DynamoDB::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::Dynamodb
+
+        %w[create_table
+          delete_item
+          delete_table
+          get_item
+          put_item
+          query
+          scan
+          update_item].each do |method_name|
+          alias_method("#{method_name}_without_new_relic".to_sym, method_name.to_sym)
+
+          define_method(method_name) do |*args|
+            instrument_method_with_new_relic(method_name, *args) { send("#{method_name}_without_new_relic".to_sym, *args) }
+          end
+        end
+
+        alias_method(:build_request_without_new_relic, :build_request)
+
+        def build_request(*args)
+          build_request_with_new_relic(*args) { build_request_without_new_relic(*args) }
+        end
+      end
+    end
+  end
+end

lib/new_relic/agent/instrumentation/dynamodb/instrumentation.rb

@@ -0,0 +1,47 @@
+# 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 Dynamodb
+    PRODUCT = 'DynamoDB'
+    DEFAULT_HOST = 'dynamodb.amazonaws.com'
+
+    def instrument_method_with_new_relic(method_name, *args)
+      return yield unless NewRelic::Agent::Tracer.tracing_enabled?
+
+      NewRelic::Agent.record_instrumentation_invocation(PRODUCT)
+
+      segment = NewRelic::Agent::Tracer.start_datastore_segment(
+        product: PRODUCT,
+        operation: method_name,
+        host: config&.endpoint&.host || DEFAULT_HOST,
+        port_path_or_id: config&.endpoint&.port,
+        collection: args[0][:table_name]
+      )
+
+      arn = get_arn(args[0])
+      segment.add_agent_attribute('cloud.resource_id', arn) if arn
+
+      @nr_captured_request = nil # clear request just in case
+      begin
+        NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+      ensure
+        segment.add_agent_attribute('aws.operation', method_name)
+        segment.add_agent_attribute('aws.requestId', @nr_captured_request&.context&.http_response&.headers&.[]('x-amzn-requestid'))
+        segment.add_agent_attribute('aws.region', config&.region)
+        segment&.finish
+      end
+    end
+
+    def build_request_with_new_relic(*args)
+      @nr_captured_request = yield
+    end
+
+    def get_arn(params)
+      return unless params[:table_name]
+
+      NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config)
+    end
+  end
+end

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

@@ -0,0 +1,26 @@
+# 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 Dynamodb::Prepend
+    include NewRelic::Agent::Instrumentation::Dynamodb
+
+    %w[create_table
+      delete_item
+      delete_table
+      get_item
+      put_item
+      query
+      scan
+      update_item].each do |method_name|
+        define_method(method_name) do |*args|
+          instrument_method_with_new_relic(method_name, *args) { super(*args) }
+        end
+      end
+
+    def build_request(*args)
+      build_request_with_new_relic(*args) { super }
+    end
+  end
+end

lib/new_relic/agent/instrumentation/grpc/client/instrumentation.rb

@@ -50,7 +50,6 @@ def add_attributes(segment, attributes_hash)
             attributes_hash.each do |attr, value|
               segment.add_agent_attribute(attr, value)
             end
-            segment.record_agent_attributes = true
           end
 
           def grpc_status_and_message_from_exception(exception)

lib/new_relic/agent/span_event_primitive.rb

@@ -79,17 +79,13 @@ def for_external_request_segment(segment)
         intrinsics[SPAN_KIND_KEY] = CLIENT
         intrinsics[SERVER_ADDRESS_KEY] = segment.uri.host
         intrinsics[SERVER_PORT_KEY] = segment.uri.port
-        agent_attributes = error_attributes(segment) || {}
+        agent_attributes = {}
 
         if allowed?(HTTP_URL_KEY)
           agent_attributes[HTTP_URL_KEY] = truncate(segment.uri)
         end
 
-        if segment.respond_to?(:record_agent_attributes?) && segment.record_agent_attributes?
-          agent_attributes.merge!(agent_attributes(segment))
-        end
-
-        [intrinsics, custom_attributes(segment), agent_attributes]
+        [intrinsics, custom_attributes(segment), agent_attributes.merge(agent_attributes(segment))]
       end
 
       def for_datastore_segment(segment) # rubocop:disable Metrics/AbcSize
@@ -99,7 +95,7 @@ def for_datastore_segment(segment) # rubocop:disable Metrics/AbcSize
         intrinsics[SPAN_KIND_KEY] = CLIENT
         intrinsics[CATEGORY_KEY] = DATASTORE_CATEGORY
 
-        agent_attributes = error_attributes(segment) || {}
+        agent_attributes = {}
 
         if segment.database_name && allowed?(DB_INSTANCE_KEY)
           agent_attributes[DB_INSTANCE_KEY] = truncate(segment.database_name)
@@ -123,7 +119,7 @@ def for_datastore_segment(segment) # rubocop:disable Metrics/AbcSize
           agent_attributes[DB_STATEMENT_KEY] = truncate(segment.nosql_statement, 2000)
         end
 
-        [intrinsics, custom_attributes(segment), agent_attributes]
+        [intrinsics, custom_attributes(segment), agent_attributes.merge(agent_attributes(segment))]
       end
 
       private

lib/new_relic/agent/transaction/external_request_segment.rb

@@ -23,7 +23,6 @@ class ExternalRequestSegment < Segment
         MISSING_STATUS_CODE = 'MissingHTTPStatusCode'
 
         attr_reader :library, :uri, :procedure, :http_status_code
-        attr_writer :record_agent_attributes
 
         def initialize(library, uri, procedure, start_time = nil) # :nodoc:
           @library = library
@@ -32,7 +31,6 @@ def initialize(library, uri, procedure, start_time = nil) # :nodoc:
           @host_header = nil
           @app_data = nil
           @http_status_code = nil
-          @record_agent_attributes = false
           super(nil, nil, start_time)
         end
 
@@ -44,14 +42,6 @@ def host # :nodoc:
           @host_header || uri.host
         end
 
-        # By default external request segments only have errors and the http
-        # url recorded as agent attributes. To have all the agent attributes
-        # recorded, use the attr_writer like so `segment.record_agent_attributes = true`
-        # See: SpanEventPrimitive#for_external_request_segment
-        def record_agent_attributes?
-          @record_agent_attributes
-        end
-
         # This method adds New Relic request headers to a given request made to an
         # external API and checks to see if a host header is used for the request.
         # If a host header is used, it updates the segment name to match the host

newrelic.yml

@@ -403,6 +403,10 @@ common: &default_settings
   # prepend, chain, disabled.
   # instrumentation.bunny: auto
 
+  # Controls auto-instrumentation of DynamoDb at start-up.
+  # May be one of [auto|prepend|chain|disabled]
+  # instrumentation.dynamodb: auto
+
   # Controls auto-instrumentation of the concurrent-ruby library at start-up. May be
   # one of: auto, prepend, chain, disabled.
   # instrumentation.concurrent_ruby: auto

test/multiverse/suites/dynamodb/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-dynamodb'
+  gem 'nokogiri'
+RB

test/multiverse/suites/dynamodb/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:
+    dynamodb: <%= $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