Kinesis instrumentation

CHANGELOG.md

@@ -6,6 +6,10 @@
 
   The agent now supports Ruby 3.4.0. We've made incremental changes throughout the preview stage to reach compatibility. This release includes an update to the Thread Profiler for compatibility with Ruby 3.4.0's new backtrace format. [Issue#2992](https://github.com/newrelic/newrelic-ruby-agent/issues/2992) [PR#2997](https://github.com/newrelic/newrelic-ruby-agent/pull/2997)
 
+- **Feature: Add instrumentation for aws-sdk-kinesis**
+
+  The agent now has instrumentation for the [aws-sdk-kinesis](https://rubygems.org/gems/aws-sdk-kinesis) gem. It will record message broker segments for `get_records`, `put_record`, and `put_records` operations. All other operations will record standard segments. [PR#2974](https://github.com/newrelic/newrelic-ruby-agent/pull/2974)
+
 - **Bugfix: Stop emitting inaccurate debug-level log about deprecated configuration options**
 
   In the previous major release, we dropped support for `disable_<library_name>` configuration options in favor of `instrumentation.<library_name>`. Previously, a DEBUG level log warning appeared whenever `disable_*` options were set to `true`, even for libraries (e.g. Action Dispatch) without equivalent `instrumentation.*` options:

lib/new_relic/agent/configuration/default_source.rb

@@ -1562,6 +1562,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.aws_sdk_kinesis' => {
+          :default => 'auto',
+          :documentation_default => 'auto',
+          :public => true,
+          :type => String,
+          :dynamic_name => true,
+          :allowed_from_server => false,
+          :description => 'Controls auto-instrumentation of the aws-sdk-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/aws_sdk_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 'aws_sdk_kinesis/instrumentation'
+require_relative 'aws_sdk_kinesis/chain'
+require_relative 'aws_sdk_kinesis/prepend'
+
+DependencyDetection.defer do
+  named :aws_sdk_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/aws_sdk_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/aws_sdk_kinesis/instrumentation.rb

@@ -0,0 +1,91 @@
+# 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'
+    MESSAGE_BROKER_SEGMENT_METHODS = %w[put_record put_records get_records].freeze
+
+    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]
+      arn = get_arn(params) if params
+
+      if MESSAGE_BROKER_SEGMENT_METHODS.include?(method_name)
+        stream_name = get_stream_name(params, arn)
+        segment = NewRelic::Agent::Tracer.start_message_broker_segment(
+          action: method_name == 'get_records' ? :consume : :produce,
+          library: KINESIS,
+          destination_type: :stream,
+          destination_name: stream_name
+        )
+      else
+        segment = NewRelic::Agent::Tracer.start_segment(name: get_segment_name(method_name, params))
+      end
+
+      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&.finish
+      end
+    end
+
+    def get_segment_name(method_name, params)
+      stream_name = params&.dig(:stream_name)
+      return "#{KINESIS}/#{method_name}/#{stream_name}" if stream_name
+
+      "#{KINESIS}/#{method_name}"
+    rescue => e
+      NewRelic::Agent.logger.warn("Failed to create segment name: #{e}")
+    end
+
+    def get_stream_name(params, arn)
+      params&.dig(:stream_name) || arn.split('/').last || 'unknown'
+    rescue => e
+      NewRelic::Agent.logger.warn("Failed to get stream 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)
+      stream_arn = params&.dig(:stream_arn)
+      return stream_arn if stream_arn
+
+      stream_name = params&.dig(:stream_name)
+      NewRelic::Agent::Aws.create_arn(KINESIS.downcase, "stream/#{stream_name}", config&.region, nr_account_id) if stream_name
+    end
+  end
+end

lib/new_relic/agent/instrumentation/aws_sdk_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

lib/new_relic/agent/transaction/message_broker_segment.rb

@@ -15,13 +15,15 @@ class MessageBrokerSegment < Segment
         PRODUCE = 'Produce'.freeze
         QUEUE = 'Queue'.freeze
         PURGE = 'Purge'.freeze
+        STREAM = 'Stream'.freeze
         TEMP = 'Temp'.freeze
         TOPIC = 'Topic'.freeze
         UNKNOWN = 'Unknown'.freeze
 
         DESTINATION_TYPES = [
           :exchange,
           :queue,
+          :stream,
           :topic,
           :temporary_queue,
           :temporary_topic,
@@ -37,6 +39,7 @@ class MessageBrokerSegment < Segment
         TYPES = {
           exchange: EXCHANGE,
           temporary_queue: QUEUE,
+          stream: STREAM,
           queue: QUEUE,
           temporary_topic: TOPIC,
           topic: TOPIC,

test/multiverse/suites/aws_sdk_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/aws_sdk_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:
+    aws_sdk_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