Kinesis instrumentation

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## dev
 
+- **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. [PR#2974](https://github.com/newrelic/newrelic-ruby-agent/pull/2974)
+
 - **Bugfix: Do not attempt to decorate logs with `nil` messages**
 
   The agent no longer attempts to add New Relic linking metadata to logs with `nil` messages. Thank you, [@arlando](https://github.com/arlando) for bringing this to our attention! [Issue#2985](https://github.com/newrelic/newrelic-ruby-agent/issues/2985) [PR#2986](https://github.com/newrelic/newrelic-ruby-agent/pull/2986)

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.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,88 @@
+# 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)
+      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 get_stream_name(params, arn)
+      params[: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)
+      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/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