newrelic_rpm 9.12.0 → 9.17.0

data/lib/new_relic/agent/instrumentation/aws_sdk_lambda/chain.rb

@@ -0,0 +1,33 @@
+# 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 'instrumentation'
+
+module NewRelic::Agent::Instrumentation
+  module AwsSdkLambda::Chain
+    def self.instrument!
+      ::Aws::Lambda::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::AwsSdkLambda
+
+        alias_method(:invoke_without_new_relic, :invoke)
+
+        def invoke(*args)
+          invoke_with_new_relic(*args) { invoke_without_new_relic(*args) }
+        end
+
+        alias_method(:invoke_async_without_new_relic, :invoke_async)
+
+        def invoke_async(*args)
+          invoke_async_with_new_relic(*args) { invoke_async_without_new_relic(*args) }
+        end
+
+        alias_method(:invoke_with_response_stream_without_new_relic, :invoke_with_response_stream)
+
+        def invoke_with_response_stream(*args)
+          invoke_with_response_stream_with_new_relic(*args) { invoke_with_response_stream_without_new_relic(*args) }
+        end
+      end
+    end
+  end
+end

data/lib/new_relic/agent/instrumentation/aws_sdk_lambda/instrumentation.rb

@@ -0,0 +1,93 @@
+# 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 'json'
+
+module NewRelic::Agent::Instrumentation
+  module AwsSdkLambda
+    INSTRUMENTATION_NAME = 'aws_sdk_lambda'
+    AWS_SERVICE = 'lambda'
+    CLOUD_PLATFORM = 'aws_lambda'
+    WRAPPED_RESPONSE = Struct.new(:status_code, :has_status_code?)
+
+    def invoke_with_new_relic(*args)
+      with_tracing(:invoke, *args) { yield }
+    end
+
+    def invoke_async_with_new_relic(*args)
+      with_tracing(:invoke_async, *args) { yield }
+    end
+
+    def invoke_with_response_stream_with_new_relic(*args)
+      with_tracing(:invoke_with_response_stream, *args) { yield }
+    end
+
+    private
+
+    def with_tracing(action, *args)
+      segment = generate_segment(action, *args)
+
+      # prevent additional instrumentation for things like Net::HTTP from
+      # creating any segments that may appear as redundant / confusing
+      NewRelic::Agent.disable_all_tracing do
+        response = NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+        process_response(response, segment)
+        response
+      end
+    ensure
+      segment&.finish
+    end
+
+    def process_response(response, segment)
+      process_function_error(response) if response.respond_to?(:function_error)
+    rescue => e
+      NewRelic::Agent.logger.error("Error processing aws-sdk-lambda invocation response: #{e}")
+    end
+
+    # notice error that was raised / unhandled by the function
+    def process_function_error(response)
+      function_error = response.function_error
+      return unless function_error
+
+      msg = "[#{function_error}]"
+      payload = response.payload&.string if response.respond_to?(:payload)
+      payload_hash = JSON.parse(payload) if payload
+      msg = "#{msg} #{payload_hash['errorType']} - #{payload_hash['errorMessage']}" if payload_hash
+      e = StandardError.new(msg)
+      e.set_backtrace(payload_hash['stackTrace']) if payload_hash
+
+      NewRelic::Agent.notice_error(e)
+    end
+
+    def generate_segment(action, options = {})
+      function = function_name(options)
+      region = aws_region
+      arn = aws_arn(function, region)
+      segment = NewRelic::Agent::Tracer.start_segment(name: "Lambda/#{action}/#{function}")
+      segment.add_agent_attribute('cloud.account.id', nr_account_id)
+      segment.add_agent_attribute('cloud.platform', CLOUD_PLATFORM)
+      segment.add_agent_attribute('cloud.region', region)
+      segment.add_agent_attribute('cloud.resource_id', arn) if arn
+      segment
+    end
+
+    def function_name(options = {})
+      (options.fetch(:function_name, nil) if options.respond_to?(:fetch)) || NewRelic::UNKNOWN
+    end
+
+    def aws_region
+      config&.region if self.respond_to?(:config)
+    end
+
+    def aws_arn(function, region)
+      NewRelic::Agent::Aws.create_arn(AWS_SERVICE, "function:#{function}", region, nr_account_id)
+    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
+  end
+end

data/lib/new_relic/agent/instrumentation/aws_sdk_lambda/prepend.rb

@@ -0,0 +1,23 @@
+# 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 'instrumentation'
+
+module NewRelic::Agent::Instrumentation
+  module AwsSdkLambda::Prepend
+    include NewRelic::Agent::Instrumentation::AwsSdkLambda
+
+    def invoke(*args)
+      invoke_with_new_relic(*args) { super }
+    end
+
+    def invoke_async(*args)
+      invoke_async_with_new_relic(*args) { super }
+    end
+
+    def invoke_with_response_stream(*args)
+      invoke_with_response_stream_with_new_relic(*args) { super }
+    end
+  end
+end

data/lib/new_relic/agent/instrumentation/aws_sdk_lambda.rb

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

data/lib/new_relic/agent/instrumentation/aws_sqs.rb

@@ -14,8 +14,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing aws-sdk-sqs instrumentation')
-
     if use_prepend?
       prepend_instrument Aws::SQS::Client, NewRelic::Agent::Instrumentation::AwsSqs::Prepend
     else

data/lib/new_relic/agent/instrumentation/bunny.rb

@@ -14,7 +14,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing Bunny instrumentation')
     require 'new_relic/agent/distributed_tracing/cross_app_tracing'
     require 'new_relic/agent/messaging'
     require 'new_relic/agent/transaction/message_broker_segment'
@@ -22,9 +21,9 @@ DependencyDetection.defer do
 
   executes do
     if use_prepend?
-      prepend_instrument Bunny::Exchange, NewRelic::Agent::Instrumentation::Bunny::Prepend::Exchange
-      prepend_instrument Bunny::Queue, NewRelic::Agent::Instrumentation::Bunny::Prepend::Queue
-      prepend_instrument Bunny::Consumer, NewRelic::Agent::Instrumentation::Bunny::Prepend::Consumer
+      prepend_instrument Bunny::Exchange, NewRelic::Agent::Instrumentation::Bunny::Prepend::Exchange, 'Bunny::Exchange'
+      prepend_instrument Bunny::Queue, NewRelic::Agent::Instrumentation::Bunny::Prepend::Queue, 'Bunny::Queue'
+      prepend_instrument Bunny::Consumer, NewRelic::Agent::Instrumentation::Bunny::Prepend::Consumer, 'Bunny::Consumer'
     else
       chain_instrument NewRelic::Agent::Instrumentation::Bunny::Chain
     end

data/lib/new_relic/agent/instrumentation/concurrent_ruby.rb

@@ -16,8 +16,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing concurrent-ruby instrumentation')
-
     if use_prepend?
       prepend_instrument(Concurrent::ThreadPoolExecutor, NewRelic::Agent::Instrumentation::ConcurrentRuby::Prepend)
 

data/lib/new_relic/agent/instrumentation/curb.rb

@@ -16,17 +16,16 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing Curb instrumentation')
     require 'new_relic/agent/distributed_tracing/cross_app_tracing'
     require 'new_relic/agent/http_clients/curb_wrappers'
   end
 
   executes do
     if use_prepend?
-      prepend_instrument Curl::Easy, NewRelic::Agent::Instrumentation::Curb::Easy::Prepend
-      prepend_instrument Curl::Multi, NewRelic::Agent::Instrumentation::Curb::Multi::Prepend
+      prepend_instrument Curl::Easy, NewRelic::Agent::Instrumentation::Curb::Easy::Prepend, 'Curb::Easy'
+      prepend_instrument Curl::Multi, NewRelic::Agent::Instrumentation::Curb::Multi::Prepend, 'Curb::Multi'
     else
-      chain_instrument NewRelic::Agent::Instrumentation::Curb::Chain
+      chain_instrument NewRelic::Agent::Instrumentation::Curb::Chain, NewRelic::Agent::Instrumentation::Curb::Multi::INSTRUMENTATION_NAME
     end
   end
 end

data/lib/new_relic/agent/instrumentation/delayed_job_instrumentation.rb

@@ -82,10 +82,6 @@ DependencyDetection.defer do
     defined?(Delayed) && defined?(Delayed::Worker)
   end
 
-  executes do
-    NewRelic::Agent.logger.info('Installing DelayedJob instrumentation [part 1/2]')
-  end
-
   executes do
     if use_prepend?
       prepend_instrument Delayed::Worker, NewRelic::Agent::Instrumentation::DelayedJob::Prepend
@@ -93,23 +89,4 @@ DependencyDetection.defer do
       chain_instrument NewRelic::Agent::Instrumentation::DelayedJob::Chain
     end
   end
-
-  executes do
-    next unless delayed_job_version < Gem::Version.new('4.1.0')
-
-    deprecation_msg = 'Instrumentation for DelayedJob versions below 4.1.0 is deprecated.' \
-      'It will stop being monitored in version 9.0.0. ' \
-      'Please upgrade your DelayedJob version to continue receiving full support. ' \
-
-    NewRelic::Agent.logger.log_once(
-      :warn,
-      :deprecated_delayed_job_version,
-      deprecation_msg
-    )
-  end
-
-  def delayed_job_version
-    # the following line needs else branch coverage
-    Gem.loaded_specs['delayed_job'].version if Gem.loaded_specs['delayed_job'] # rubocop:disable Style/SafeNavigation
-  end
 end

data/lib/new_relic/agent/instrumentation/dynamodb/instrumentation.rb

@@ -56,7 +56,7 @@ module NewRelic::Agent::Instrumentation
     end
 
     def get_arn(params)
-      return unless params[:table_name] && nr_account_id
+      return unless params[:table_name]
 
       NewRelic::Agent::Aws.create_arn(PRODUCT.downcase, "table/#{params[:table_name]}", config&.region, nr_account_id)
     end

data/lib/new_relic/agent/instrumentation/dynamodb.rb

@@ -14,8 +14,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing DynamoDB instrumentation')
-
     if use_prepend?
       prepend_instrument Aws::DynamoDB::Client, NewRelic::Agent::Instrumentation::DynamoDB::Prepend
     else

data/lib/new_relic/agent/instrumentation/elasticsearch.rb

@@ -14,8 +14,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing Elasticsearch instrumentation')
-
     to_instrument = if Gem::Version.create(Elasticsearch::VERSION) < Gem::Version.create('8.0.0')
       Elasticsearch::Transport::Client
     else

data/lib/new_relic/agent/instrumentation/ethon.rb

@@ -21,10 +21,6 @@ DependencyDetection.defer do
     defined?(Ethon) && Gem::Version.new(Ethon::VERSION) >= Gem::Version.new('0.12.0')
   end
 
-  executes do
-    NewRelic::Agent.logger.info('Installing ethon instrumentation')
-  end
-
   executes do
     if use_prepend?
       # NOTE: by default prepend_instrument will go with the module name that

data/lib/new_relic/agent/instrumentation/excon.rb

@@ -33,22 +33,6 @@ DependencyDetection.defer do
     end
   end
 
-  executes do
-    next unless Gem::Version.new(Excon::VERSION) < Gem::Version.new('0.56.0')
-
-    deprecation_msg = 'Instrumentation for Excon versions below 0.56.0 is deprecated.' \
-      'They will stop being monitored in version 9.0.0. ' \
-      'Please upgrade your Excon version to continue receiving full support. '
-
-    NewRelic::Agent.logger.log_once(
-      :warn,
-      :deprecated_excon_version,
-      deprecation_msg
-    )
-
-    NewRelic::Agent.record_metric('Supportability/Deprecated/Excon', 1)
-  end
-
   def install_excon_instrumentation(excon_version)
     require 'new_relic/agent/distributed_tracing/cross_app_tracing'
     require 'new_relic/agent/http_clients/excon_wrappers'

data/lib/new_relic/agent/instrumentation/fiber.rb

@@ -14,8 +14,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing Fiber instrumentation')
-
     if use_prepend?
       prepend_instrument Fiber, NewRelic::Agent::Instrumentation::MonitoredFiber::Prepend
     else

data/lib/new_relic/agent/instrumentation/grape/instrumentation.rb

@@ -56,10 +56,7 @@ module NewRelic::Agent::Instrumentation
 
       def name_transaction(route, class_name, version)
         txn_name = name_for_transaction(route, class_name, version)
-        segment_name = "Middleware/Grape/#{class_name}/call"
         NewRelic::Agent::Transaction.set_default_transaction_name(txn_name, :grape)
-        txn = NewRelic::Agent::Transaction.tl_current
-        txn.segments.last.name = segment_name
       end
 
       def name_for_transaction(route, class_name, version)

data/lib/new_relic/agent/instrumentation/grape.rb

@@ -19,7 +19,7 @@ DependencyDetection.defer do
 
   depends_on do
     begin
-      if defined?(Bundler) && Bundler.rubygems.all_specs.map(&:name).include?('newrelic-grape')
+      if NewRelic::Helper.rubygems_specs.map(&:name).include?('newrelic-grape')
         NewRelic::Agent.logger.info('Not installing New Relic supported Grape instrumentation because the third party newrelic-grape gem is present')
         false
       else

data/lib/new_relic/agent/instrumentation/httpclient.rb

@@ -23,7 +23,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing HTTPClient instrumentation')
     require 'new_relic/agent/distributed_tracing/cross_app_tracing'
     require 'new_relic/agent/http_clients/httpclient_wrappers'
   end

data/lib/new_relic/agent/instrumentation/httprb.rb

@@ -14,7 +14,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing http.rb Wrappers')
     require 'new_relic/agent/distributed_tracing/cross_app_tracing'
     require 'new_relic/agent/http_clients/http_rb_wrappers'
   end

data/lib/new_relic/agent/instrumentation/httpx.rb

@@ -13,10 +13,6 @@ DependencyDetection.defer do
     defined?(HTTPX) && Gem::Version.new(HTTPX::VERSION) >= Gem::Version.new('1.0.0')
   end
 
-  executes do
-    NewRelic::Agent.logger.info('Installing httpx instrumentation')
-  end
-
   executes do
     if use_prepend?
       prepend_instrument HTTPX::Session, NewRelic::Agent::Instrumentation::HTTPX::Prepend

data/lib/new_relic/agent/instrumentation/logger.rb

@@ -15,12 +15,10 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing Logger instrumentation')
-
     if use_prepend?
       prepend_instrument Logger, NewRelic::Agent::Instrumentation::Logger::Prepend
     else
-      chain_instrument NewRelic::Agent::Instrumentation::Logger
+      chain_instrument NewRelic::Agent::Instrumentation::Logger, NewRelic::Agent::Instrumentation::Logger::INSTRUMENTATION_NAME
     end
   end
 end

data/lib/new_relic/agent/instrumentation/logstasher.rb

@@ -16,8 +16,6 @@ DependencyDetection.defer do
   end
 
   executes do
-    NewRelic::Agent.logger.info('Installing LogStasher instrumentation')
-
     if use_prepend?
       prepend_instrument LogStasher.singleton_class, NewRelic::Agent::Instrumentation::LogStasher::Prepend
     else

data/lib/new_relic/agent/instrumentation/memcache.rb

@@ -74,7 +74,6 @@ DependencyDetection.defer do
   depends_on { NewRelic::Agent::Instrumentation::Memcache::DalliCAS.should_instrument? }
 
   executes do
-    NewRelic::Agent.logger.info('Installing Dalli CAS Client Memcache instrumentation')
     if use_prepend?
       prepend_module = NewRelic::Agent::Instrumentation::Memcache::Prepend
       prepend_module.dalli_cas_prependers do |client_class, instrumenting_module|

data/lib/new_relic/agent/instrumentation/opensearch/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 OpenSearch::Chain
+    def self.instrument!
+      ::OpenSearch::Transport::Client.class_eval do
+        include NewRelic::Agent::Instrumentation::OpenSearch
+
+        alias_method(:perform_request_without_tracing, :perform_request)
+
+        def perform_request(*args)
+          perform_request_with_tracing(*args) do
+            perform_request_without_tracing(*args)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/new_relic/agent/instrumentation/opensearch/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 OpenSearch
+    PRODUCT_NAME = 'OpenSearch'
+    OPERATION = 'perform_request'
+    OPERATION_PATTERN = %r{/lib/opensearch/api/(?!.+#{OPERATION})}
+    INSTANCE_METHOD_PATTERN = /:in (?:`|')(?:.+#)?([^']+)'\z/
+    INSTRUMENTATION_NAME = NewRelic::Agent.base_name(name)
+
+    def perform_request_with_tracing(_method, _path, params = NewRelic::EMPTY_HASH, body = nil, _headers = nil, _opts = nil, &_block)
+      return yield unless NewRelic::Agent::Tracer.tracing_enabled?
+
+      segment = NewRelic::Agent::Tracer.start_datastore_segment(
+        product: PRODUCT_NAME,
+        operation: nr_operation || OPERATION,
+        host: nr_hosts[:host],
+        port_path_or_id: nr_hosts[:port],
+        database_name: nr_cluster_name
+      )
+      begin
+        NewRelic::Agent::Tracer.capture_segment_error(segment) { yield }
+      ensure
+        if segment
+          segment.notice_nosql_statement(nr_reported_query(body || params))
+          segment.finish
+        end
+      end
+    end
+
+    private
+
+    # See Elasticsearch instrumentation for explanation on Ruby 3.4 changes to match instance method
+    def nr_operation
+      location = caller_locations.detect { |loc| loc.to_s.match?(OPERATION_PATTERN) }
+      return unless location && location.to_s =~ INSTANCE_METHOD_PATTERN
+
+      Regexp.last_match(1)
+    end
+
+    def nr_reported_query(query)
+      return unless NewRelic::Agent.config[:'opensearch.capture_queries']
+      return query unless NewRelic::Agent.config[:'opensearch.obfuscate_queries']
+
+      NewRelic::Agent::Datastores::NosqlObfuscator.obfuscate_statement(query)
+    end
+
+    def nr_cluster_name
+      return @nr_cluster_name if defined?(@nr_cluster_name)
+      return if nr_hosts.empty?
+
+      NewRelic::Agent.disable_all_tracing do
+        @nr_cluster_name ||= perform_request('GET', '/').body['cluster_name']
+      end
+    rescue StandardError => e
+      NewRelic::Agent.logger.error('Failed to get cluster name for OpenSearch', e)
+      nil
+    end
+
+    def nr_hosts
+      @nr_hosts ||= (transport.hosts.first || NewRelic::EMPTY_HASH)
+    end
+  end
+end

data/lib/new_relic/agent/instrumentation/opensearch/prepend.rb

@@ -0,0 +1,13 @@
+# 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 OpenSearch::Prepend
+    include NewRelic::Agent::Instrumentation::OpenSearch
+
+    def perform_request(*args)
+      perform_request_with_tracing(*args) { super }
+    end
+  end
+end

data/lib/new_relic/agent/instrumentation/opensearch.rb

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

data/lib/new_relic/agent/instrumentation/padrino.rb

@@ -22,11 +22,11 @@ DependencyDetection.defer do
   depends_on { defined?(Padrino) && defined?(Padrino::Routing::InstanceMethods) }
 
   executes do
-    NewRelic::Agent.logger.info('Installing Padrino instrumentation')
+    supportability_name = NewRelic::Agent::Instrumentation::Padrino::INSTRUMENTATION_NAME
     if use_prepend?
-      prepend_instrument Padrino::Application, NewRelic::Agent::Instrumentation::PadrinoTracer::Prepend
+      prepend_instrument Padrino::Application, NewRelic::Agent::Instrumentation::PadrinoTracer::Prepend, supportability_name
     else
-      chain_instrument NewRelic::Agent::Instrumentation::PadrinoTracer::Chain
+      chain_instrument NewRelic::Agent::Instrumentation::PadrinoTracer::Chain, supportability_name
     end
   end
 end

data/lib/new_relic/agent/instrumentation/rake.rb

@@ -17,7 +17,6 @@ DependencyDetection.defer do
   depends_on { NewRelic::Agent::Instrumentation::Rake.safe_from_third_party_gem? }
 
   executes do
-    NewRelic::Agent.logger.info('Installing Rake instrumentation')
     NewRelic::Agent.logger.debug("Instrumenting Rake tasks: #{NewRelic::Agent.config[:'rake.tasks']}")
   end
 

data/lib/new_relic/agent/instrumentation/rdkafka/chain.rb

@@ -0,0 +1,72 @@
+# 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 'instrumentation'
+
+module NewRelic::Agent::Instrumentation
+  module Rdkafka::Chain
+    def self.instrument!
+      ::Rdkafka::Producer.class_eval do
+        include NewRelic::Agent::Instrumentation::Rdkafka
+
+        alias_method(:produce_without_new_relic, :produce)
+
+        def produce(**kwargs)
+          produce_with_new_relic(kwargs) do |headers|
+            kwargs[:headers] = headers
+            produce_without_new_relic(**kwargs)
+          end
+        end
+      end
+
+      ::Rdkafka::Consumer.class_eval do
+        include NewRelic::Agent::Instrumentation::Rdkafka
+
+        alias_method(:each_without_new_relic, :each)
+
+        def each(**kwargs)
+          each_without_new_relic(**kwargs) do |message|
+            each_with_new_relic(message) do
+              yield(message)
+            end
+          end
+        end
+      end
+
+      ::Rdkafka::Config.class_eval do
+        include NewRelic::Agent::Instrumentation::RdkafkaConfig
+
+        alias_method(:producer_without_new_relic, :producer)
+        alias_method(:consumer_without_new_relic, :consumer)
+
+        if Gem::Version.new(::Rdkafka::VERSION) >= Gem::Version.new('0.16.0') ||
+            (Gem::Version.new(RUBY_VERSION) >= Gem::Version.new('2.7.0'))
+          def producer(**kwargs)
+            producer_without_new_relic(**kwargs).tap do |producer|
+              set_nr_config(producer)
+            end
+          end
+
+          def consumer(**kwargs)
+            consumer_without_new_relic(**kwargs).tap do |consumer|
+              set_nr_config(consumer)
+            end
+          end
+        else
+          def producer
+            producer_without_new_relic.tap do |producer|
+              set_nr_config(producer)
+            end
+          end
+
+          def consumer
+            consumer_without_new_relic.tap do |consumer|
+              set_nr_config(consumer)
+            end
+          end
+        end
+      end
+    end
+  end
+end