newrelic_rpm 9.9.0 → 9.21.0

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/grpc/client/instrumentation.rb

@@ -50,7 +50,6 @@ module NewRelic
             attributes_hash.each do |attr, value|
               segment.add_agent_attribute(attr, value)
             end
-            segment.record_agent_attributes = true
           end
 
           def grpc_status_and_message_from_exception(exception)

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

@@ -16,14 +16,10 @@ DependencyDetection.defer do
   end
 
   depends_on do
-    minimum_supported_version = Gem::Version.new(HTTPCLIENT_MIN_VERSION)
-    current_version = Gem::Version.new(HTTPClient::VERSION)
-
-    current_version >= minimum_supported_version
+    NewRelic::Helper.version_satisfied?(HTTPClient::VERSION, '>=', HTTPCLIENT_MIN_VERSION)
   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/instrumentation.rb

@@ -30,7 +30,7 @@ module NewRelic::Agent::Instrumentation::HTTPX
 
   def nr_finish_segment
     proc do |request, segment|
-      response = @responses[request]
+      response = request.response if request
 
       unless response
         NewRelic::Agent.logger.debug('Processed an on-response callback for HTTPX but could not find the response!')

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

@@ -10,11 +10,7 @@ DependencyDetection.defer do
   named :httpx
 
   depends_on do
-    defined?(HTTPX) && Gem::Version.new(HTTPX::VERSION) >= Gem::Version.new('1.0.0')
-  end
-
-  executes do
-    NewRelic::Agent.logger.info('Installing httpx instrumentation')
+    defined?(HTTPX) && NewRelic::Helper.version_satisfied?(HTTPX::VERSION, '>=', '1.0.0')
   end
 
   executes do

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/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 LogStasher::Chain
+    def self.instrument!
+      ::LogStasher.singleton_class.class_eval do
+        include NewRelic::Agent::Instrumentation::LogStasher
+
+        alias_method(:build_logstash_event_without_new_relic, :build_logstash_event)
+
+        def build_logstash_event(*args)
+          build_logstash_event_with_new_relic(*args) do
+            build_logstash_event_without_new_relic(*args)
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,24 @@
+# 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 LogStasher
+    INSTRUMENTATION_NAME = NewRelic::Agent.base_name(name)
+
+    def self.enabled?
+      NewRelic::Agent.config[:'instrumentation.logstasher'] != 'disabled'
+    end
+
+    def build_logstash_event_with_new_relic(*args)
+      logstasher_event = yield
+      log = logstasher_event.instance_variable_get(:@data)
+
+      ::NewRelic::Agent.record_instrumentation_invocation(INSTRUMENTATION_NAME)
+      ::NewRelic::Agent.agent.log_event_aggregator.record_logstasher_event(log)
+      ::NewRelic::Agent::LocalLogDecorator.decorate(log)
+
+      logstasher_event
+    end
+  end
+end

data/lib/new_relic/agent/instrumentation/logstasher/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 LogStasher::Prepend
+    include NewRelic::Agent::Instrumentation::LogStasher
+
+    def build_logstash_event(*args)
+      build_logstash_event_with_new_relic(*args) { super }
+    end
+  end
+end

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

@@ -0,0 +1,25 @@
+# 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 'logstasher/instrumentation'
+require_relative 'logstasher/chain'
+require_relative 'logstasher/prepend'
+
+DependencyDetection.defer do
+  named :logstasher
+
+  depends_on do
+    defined?(LogStasher) &&
+      NewRelic::Helper.version_satisfied?(LogStasher::VERSION, '>=', '1.0.0') &&
+      NewRelic::Agent.config[:'application_logging.enabled']
+  end
+
+  executes do
+    if use_prepend?
+      prepend_instrument LogStasher.singleton_class, NewRelic::Agent::Instrumentation::LogStasher::Prepend
+    else
+      chain_instrument NewRelic::Agent::Instrumentation::LogStasher::Chain
+    end
+  end
+end

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

@@ -61,7 +61,7 @@ module NewRelic
               # TODO: MAJOR VERSION
               # Dalli - 3.1.0 renamed send_multiget to pipelined_get, but the method is otherwise the same
               # Once we no longer support Dalli < 3.1.0, remove this conditional logic
-              if Gem::Version.new(::Dalli::VERSION) >= Gem::Version.new('3.1.0')
+              if NewRelic::Helper.version_satisfied?(::Dalli::VERSION, '>=', '3.1.0')
                 alias_method(:pipelined_get_without_newrelic_trace, :pipelined_get)
                 def pipelined_get(keys)
                   send_multiget_with_newrelic_tracing(keys) { pipelined_get_without_newrelic_trace(keys) }

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

@@ -9,11 +9,11 @@ module NewRelic::Agent::Instrumentation
       BINARY_PROTOCOL_SUPPORTED_VERSION = Gem::Version.new('3.0.2')
 
       def supports_datastore_instances?
-        DATASTORE_INSTANCES_SUPPORTED_VERSION <= Gem::Version.new(::Dalli::VERSION)
+        NewRelic::Helper.version_satisfied?(DATASTORE_INSTANCES_SUPPORTED_VERSION, '<=', ::Dalli::VERSION)
       end
 
       def supports_binary_protocol?
-        BINARY_PROTOCOL_SUPPORTED_VERSION <= Gem::Version.new(::Dalli::VERSION)
+        NewRelic::Helper.version_satisfied?(BINARY_PROTOCOL_SUPPORTED_VERSION, '<=', ::Dalli::VERSION)
       end
 
       def client_methods

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

@@ -6,7 +6,7 @@ module NewRelic::Agent::Instrumentation
   module Memcache
     module Tracer
       SLASH = '/'
-      UNKNOWN = 'unknown'
+      UNKNOWN = NewRelic::UNKNOWN_LOWER
       LOCALHOST = 'localhost'
       MULTIGET_METRIC_NAME = 'get_multi_request'
       MEMCACHED = 'Memcached'

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

@@ -87,7 +87,7 @@ module NewRelic::Agent::Instrumentation
           # TODO: MAJOR VERSION
           # Dalli - 3.1.0 renamed send_multiget to pipelined_get, but the method is otherwise the same
           # Once we no longer support Dalli < 3.1.0, remove this conditional logic
-          if Gem::Version.new(::Dalli::VERSION) >= Gem::Version.new('3.1.0')
+          if NewRelic::Helper.version_satisfied?(::Dalli::VERSION, '>=', '3.1.0')
             def pipelined_get(keys)
               send_multiget_with_newrelic_tracing(keys) { super }
             end

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/mongodb_command_subscriber.rb

@@ -106,7 +106,7 @@ module NewRelic
           )
         end
 
-        UNKNOWN = 'unknown'.freeze
+        UNKNOWN = NewRelic::UNKNOWN_LOWER
         LOCALHOST = 'localhost'.freeze
 
         def host_from_address(address)

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

@@ -21,7 +21,7 @@ module NewRelic
 
           begin
             response = nil
-            segment.add_request_headers(wrapped_request)
+            segment&.add_request_headers(wrapped_request)
 
             # RUBY-1244 Disable further tracing in request to avoid double
             # counting if connection wasn't started (which calls request again).
@@ -34,10 +34,10 @@ module NewRelic
             wrapped_response = NewRelic::Agent::HTTPClients::NetHTTPResponse.new(response)
 
             if NewRelic::Agent::LLM.openai_parent?(segment)
-              NewRelic::Agent::LLM.populate_openai_response_headers(wrapped_response, segment.parent)
+              NewRelic::Agent::LLM.populate_openai_response_headers(wrapped_response, segment&.parent)
             end
 
-            segment.process_response_headers(wrapped_response)
+            segment&.process_response_headers(wrapped_response)
 
             response
           ensure

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

@@ -19,7 +19,8 @@ DependencyDetection.defer do
 
   # Airbrake uses method chaining on Net::HTTP in versions < 10.0.2 (10.0.2 updated to prepend for Net:HTTP)
   conflicts_with_prepend do
-    defined?(Airbrake) && defined?(Airbrake::AIRBRAKE_VERSION) && Gem::Version.create(Airbrake::AIRBRAKE_VERSION) < Gem::Version.create('10.0.2')
+    defined?(Airbrake) && defined?(Airbrake::AIRBRAKE_VERSION) &&
+      NewRelic::Helper.version_satisfied?(Airbrake::AIRBRAKE_VERSION, '<', '10.0.2')
   end
 
   conflicts_with_prepend do

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

@@ -16,8 +16,6 @@ module NewRelic
         end
 
         def self.find_all_subscribers
-          # TODO: need to talk to Rails core about an API for this,
-          # rather than digging through Listener ivars
           instance_variable_names = [:@subscribers, :@string_subscribers, :@other_subscribers]
           all_subscribers = []
 

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/{tasks/instrumentation_generator/templates/instrumentation.tt → new_relic/agent/instrumentation/opensearch/prepend.rb}

@@ -3,11 +3,11 @@
 # frozen_string_literal: true
 
 module NewRelic::Agent::Instrumentation
-  module <%= @class_name %>
+  module OpenSearch::Prepend
+    include NewRelic::Agent::Instrumentation::OpenSearch
 
-    def <%= @method.downcase %>_with_new_relic<%= "(#{@args})" unless @args.empty? %>
-      # add instrumentation content here
-      yield
+    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/rack/instrumentation.rb

@@ -13,6 +13,7 @@ module NewRelic
             attr_accessor :_nr_deferred_detection_ran
           end
           builder_class._nr_deferred_detection_ran = false
+          NewRelic::Control::SecurityInterface.instance.wait = true
         end
 
         def deferred_dependency_check
@@ -21,6 +22,8 @@ module NewRelic
           NewRelic::Agent.logger.info('Doing deferred dependency-detection before Rack startup')
           DependencyDetection.detect!
           self.class._nr_deferred_detection_ran = true
+          NewRelic::Control::SecurityInterface.instance.wait = false
+          NewRelic::Control::SecurityInterface.instance.init_agent
         end
 
         def check_for_late_instrumentation(app)

data/lib/new_relic/agent/instrumentation/rails_notifications/action_controller.rb

@@ -32,15 +32,19 @@ DependencyDetection.defer do
     NewRelic::Agent::Instrumentation::ActionControllerSubscriber \
       .subscribe(/^process_action.action_controller$/)
 
-    subs = %w[send_file
+    subs = %w[exist_fragment?
+      expire_fragment
+      halted_callback
+      read_fragment
+      redirect_to
       send_data
+      send_file
       send_stream
-      redirect_to
-      halted_callback
-      unpermitted_parameters]
+      write_fragment
+      unpermitted_parameters].map { |s| Regexp.escape(s) }
 
     # have to double escape period because its going from string -> regex
     NewRelic::Agent::Instrumentation::ActionControllerOtherSubscriber \
-      .subscribe(Regexp.new("^(#{subs.join('|')})\\.action_controller$"))
+      .subscribe(Regexp.new("^(?:#{subs.join('|')})\\.action_controller$"))
   end
 end

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

@@ -12,12 +12,11 @@ DependencyDetection.defer do
   configure_with :rake
 
   depends_on { defined?(Rake) && defined?(Rake::VERSION) }
-  depends_on { Gem::Version.new(Rake::VERSION) >= Gem::Version.new('10.0.0') }
+  depends_on { NewRelic::Helper.version_satisfied?(Rake::VERSION, '>=', '10.0.0') }
   depends_on { NewRelic::Agent.config[:'rake.tasks'].any? }
   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 NewRelic::Helper.version_satisfied?(::Rdkafka::VERSION, '>=', '0.16.0') ||
+            NewRelic::Helper.version_satisfied?(RUBY_VERSION, '>=', '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

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

@@ -0,0 +1,70 @@
+# 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 'new_relic/agent/messaging'
+
+module NewRelic::Agent::Instrumentation
+  module Rdkafka
+    MESSAGING_LIBRARY = 'Kafka'
+    PRODUCE = 'Produce'
+    CONSUME = 'Consume'
+
+    INSTRUMENTATION_NAME = 'Rdkafka'
+
+    def produce_with_new_relic(*args)
+      NewRelic::Agent.record_instrumentation_invocation(INSTRUMENTATION_NAME)
+
+      topic_name = args[0][:topic]
+      segment = NewRelic::Agent::Tracer.start_message_broker_segment(
+        action: :produce,
+        library: MESSAGING_LIBRARY,
+        destination_type: :topic,
+        destination_name: topic_name
+      )
+      create_kafka_metrics(action: PRODUCE, topic: topic_name)
+
+      headers = args[0][:headers] || {}
+      ::NewRelic::Agent::DistributedTracing.insert_distributed_trace_headers(headers)
+
+      NewRelic::Agent::Tracer.capture_segment_error(segment) { yield(headers) }
+    ensure
+      segment&.finish
+    end
+
+    def each_with_new_relic(message)
+      NewRelic::Agent.record_instrumentation_invocation(INSTRUMENTATION_NAME)
+
+      headers = message&.headers || {}
+      topic_name = message&.topic
+
+      NewRelic::Agent::Messaging.wrap_message_broker_consume_transaction(
+        library: MESSAGING_LIBRARY,
+        destination_type: :topic,
+        destination_name: topic_name,
+        headers: headers,
+        action: :consume
+      ) do
+        create_kafka_metrics(action: CONSUME, topic: topic_name)
+        yield
+      end
+    end
+
+    def create_kafka_metrics(action:, topic:)
+      hosts = []
+      # both 'bootstrap.servers' and 'metadata.broker.list' are valid ways to specify the Kafka server
+      hosts << @nr_config[:'bootstrap.servers'] if @nr_config[:'bootstrap.servers']
+      hosts << @nr_config[:'metadata.broker.list'] if @nr_config[:'metadata.broker.list']
+      hosts.each do |host|
+        NewRelic::Agent.record_metric("MessageBroker/Kafka/Nodes/#{host}/#{action}/#{topic}", 1)
+        NewRelic::Agent.record_metric("MessageBroker/Kafka/Nodes/#{host}", 1)
+      end
+    end
+  end
+
+  module RdkafkaConfig
+    def set_nr_config(producer_or_consumer)
+      producer_or_consumer.instance_variable_set(:@nr_config, self)
+    end
+  end
+end