Merge pull request #2973 from newrelic/firehose_instrumentation

Firehose instrumentation

CHANGELOG.md

@@ -5,6 +5,10 @@
 - **Feature: Support Ruby 3.4.0**
 
   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-firehose**
+
+  The agent now has instrumentation for the [aws-sdk-firehose](https://rubygems.org/gems/aws-sdk-firehose) gem. [PR#2973](https://github.com/newrelic/newrelic-ruby-agent/pull/2973)
 
 - **Feature: Kubernetes APM auto-attach - new agent version precedent**
 

lib/new_relic/agent/configuration/default_source.rb

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

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

lib/new_relic/agent/instrumentation/aws_sdk_firehose/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 Firehose::Chain
+    def self.instrument!
+      ::Aws::Firehose::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::Firehose
+
+        NewRelic::Agent::Instrumentation::Firehose::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_firehose/instrumentation.rb

@@ -0,0 +1,66 @@
+# 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 Firehose
+    INSTRUMENTED_METHODS = %w[
+      create_delivery_stream
+      delete_delivery_stream
+      describe_delivery_stream
+      list_delivery_streams
+      list_tags_for_delivery_stream
+      put_record
+      put_record_batch
+      start_delivery_stream_encryption
+      stop_delivery_stream_encryption
+      tag_delivery_stream
+      untag_delivery_stream
+      update_destination
+    ].freeze
+
+    FIREHOSE = 'Firehose'
+    AWS_KINESIS_DELIVERY_STREAMS = 'aws_kinesis_delivery_streams'
+
+    def instrument_method_with_new_relic(method_name, *args)
+      return yield unless NewRelic::Agent::Tracer.tracing_enabled?
+
+      NewRelic::Agent.record_instrumentation_invocation(FIREHOSE)
+
+      params = args[0]
+      segment = NewRelic::Agent::Tracer.start_segment(name: get_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_DELIVERY_STREAMS)
+        segment&.finish
+      end
+    end
+
+    def get_segment_name(method_name, params)
+      stream_name = params&.dig(:delivery_stream_name)
+      return "#{FIREHOSE}/#{method_name}/#{stream_name}" if stream_name
+
+      "#{FIREHOSE}/#{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)
+      stream_arn = params&.dig(:delivery_stream_arn)
+      return stream_arn if stream_arn
+
+      stream_name = params&.dig(:delivery_stream_name)
+      NewRelic::Agent::Aws.create_arn(FIREHOSE.downcase, "deliverystream/#{stream_name}", config&.region, nr_account_id) if stream_name
+    end
+  end
+end

lib/new_relic/agent/instrumentation/aws_sdk_firehose/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 Firehose::Prepend
+    include NewRelic::Agent::Instrumentation::Firehose
+
+    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/aws_sdk_firehose/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-firehose'
+RB

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