Kinesis instrumentation

lib/new_relic/agent/configuration/default_source.rb

@@ -1561,6 +1561,15 @@ def self.notify
           :allowed_from_server => false,
           :description => 'Controls auto-instrumentation of the aws_sdk_lambda library at start-up. May be one of `auto`, `prepend`, `chain`, `disabled`.'
         },
+        :'instrumentation.kinesis' => {
+          :default => 'auto',
+          :documentation_default => 'auto',
+          :public => true,
+          :type => String,
+          :dynamic_name => true,
+          :allowed_from_server => false,
+          :description => 'Controls auto-instrumentation of the kinesis library at start-up. May be one of `auto`, `prepend`, `chain`, `disabled`.'
+        },
         :'instrumentation.ruby_kafka' => {
           :default => 'auto',
           :public => true,

lib/new_relic/agent/instrumentation/kinesis.rb

@@ -0,0 +1,22 @@
+# 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 'kinesis/instrumentation'
+require_relative 'kinesis/chain'
+require_relative 'kinesis/prepend'
+
+DependencyDetection.defer do
+  named :kinesis
+
+  depends_on do
+    defined?(Aws::Kinesis::Client)
+  end
+  executes do
+    if use_prepend?
+      prepend_instrument Aws::Kinesis::Client, NewRelic::Agent::Instrumentation::Kinesis::Prepend
+    else
+      chain_instrument NewRelic::Agent::Instrumentation::Kinesis::Chain
+    end
+  end
+end

lib/new_relic/agent/instrumentation/kinesis/chain.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 Kinesis::Chain
+    def self.instrument!
+      ::Aws::Kinesis::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::Kinesis
+
+        NewRelic::Agent::Instrumentation::Kinesis::INSTRUMENTED_METHODS.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
+      end
+    end
+  end
+end

lib/new_relic/agent/instrumentation/kinesis/instrumentation.rb

@@ -0,0 +1,71 @@
+# 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 Kinesis
+    INSTRUMENTED_METHODS = %w[
+      add_tags_to_stream
+      create_stream
+      decrease_stream_retention_period
+      delete_stream
+      describe_limits
+      describe_stream
+      disable_enhanced_monitoring
+      enable_enhanced_monitoring
+      get_records
+      get_shard_iterator
+      increase_stream_retention_period
+      list_streams
+      list_tags_for_stream
+      merge_shards
+      put_record
+      put_records
+      remove_tags_from_stream
+      split_shard
+      update_shard_count
+    ].freeze
+
+    KINESIS = 'Kinesis'
+    AWS_KINESIS_DATA_STREAMS = 'aws_kinesis_data_streams'
+
+    def instrument_method_with_new_relic(method_name, *args)
+      return yield unless NewRelic::Agent::Tracer.tracing_enabled?
+
+      NewRelic::Agent.record_instrumentation_invocation(KINESIS)
+
+      params = args[0]
+      segment = NewRelic::Agent::Tracer.start_segment(name: segment_name(method_name, params))
+      arn = get_arn(params) if params
+      segment&.add_agent_attribute('cloud.resource_id', arn) if arn
+
+      begin
+        NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+      ensure
+        segment&.add_agent_attribute('cloud.platform', AWS_KINESIS_DATA_STREAMS)
+        segment&.add_agent_attribute('name', segment_name(method_name, params))
+        segment&.finish
+      end
+    end
+
+    def segment_name(method_name, params)
+      return "#{KINESIS}/#{method_name}/#{params[:stream_name]}" if params&.dig(:stream_name)
+
+      "#{KINESIS}/#{method_name}"
+    rescue => e
+      NewRelic::Agent.logger.warn("Failed to create segment name: #{e}")
+    end
+
+    def nr_account_id
+      return @nr_account_id if defined?(@nr_account_id)
+
+      @nr_account_id = NewRelic::Agent::Aws.get_account_id(config)
+    end
+
+    def get_arn(params)
+      return params[:stream_arn] if params[:stream_arn]
+
+      NewRelic::Agent::Aws.create_arn(KINESIS.downcase, "stream/#{params[:stream_name]}", config&.region, nr_account_id) if params[:stream_name]
+    end
+  end
+end

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

@@ -0,0 +1,15 @@
+# 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 Kinesis::Prepend
+    include NewRelic::Agent::Instrumentation::Kinesis
+
+    INSTRUMENTED_METHODS.each do |method_name|
+      define_method(method_name) do |*args|
+        instrument_method_with_new_relic(method_name, *args) { super(*args) }
+      end
+    end
+  end
+end

test/multiverse/suites/kinesis/Envfile

@@ -0,0 +1,9 @@
+# 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-kinesis'
+RB

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