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,55 @@
+# 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]
+
+      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.debug("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('7fffffffff80', 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

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