Merge pull request #2178 from newrelic/segment_callbacks

support user defined segment creation callbacks

lib/new_relic/agent/transaction/abstract_segment.rb

@@ -24,6 +24,9 @@ class AbstractSegment
         attr_writer :record_metrics, :record_scoped_metric, :record_on_finish
         attr_reader :noticed_error
 
+        CALLBACK = :@callback
+        SEGMENT = 'segment'
+
         def initialize(name = nil, start_time = nil)
           @name = name
           @starting_segment_key = NewRelic::Agent::Tracer.current_segment_key
@@ -49,6 +52,7 @@ def initialize(name = nil, start_time = nil)
           @code_function = nil
           @code_lineno = nil
           @code_namespace = nil
+          invoke_callback
         end
 
         def start
@@ -327,6 +331,54 @@ def transaction_state
             Tracer.state
           end
         end
+
+        # for segment callback usage info, see self.set_segment_callback
+        def invoke_callback
+          return unless self.class.instance_variable_defined?(CALLBACK)
+
+          NewRelic::Agent.logger.debug("Invoking callback for #{self.class.name}...")
+          self.class.instance_variable_get(CALLBACK).call
+        end
+
+        # Setting and invoking a segment callback
+        # =======================================
+        # Each individual segment class such as `ExternalRequestSegment` allows
+        # for exactly one instance of a `Proc` (meaning a proc or lambda) to be
+        # set as a callback. A callback can be set on a segment class by calling
+        # `.set_segment_callback` with a proc or lambda as the only argument.
+        # If set, the callback will be invoked with `#call` at segment class
+        # initialization time.
+        #
+        # Example usage:
+        #   callback = -> { puts 'Hello, World! }
+        #   ExternalRequestSegment.set_segment_callback(callback)
+        #   ExternalRequestSegment.new(library, uri, procedure)
+        #
+        # A callback set on a segment class will only be called when that
+        # specific segment class is initialized. Other segment classes will not
+        # be impacted.
+        #
+        # Great caution should be taken in the defining of the callback block
+        # to not have the block perform anything too time consuming or resource
+        # intensive in order to keep the New Relic Ruby agent operating
+        # normally.
+        #
+        # Given that callbacks are user defined, they must be set entirely at
+        # the user's own risk. It is recommended that each callback use
+        # conditional logic that only performs work for certain qualified
+        # segments. It is recommended that each callback be thoroughly tested
+        # in non-production environments before being introduced to production
+        # environments.
+        def self.set_segment_callback(callback_proc)
+          unless callback_proc.is_a?(Proc)
+            NewRelic::Agent.logger.error("#{self}.#{__method__}: expected an argument of type Proc, " \
+                                         "got #{callback_proc.class}")
+            return
+          end
+
+          NewRelic::Agent.record_api_supportability_metric(:set_segment_callback)
+          instance_variable_set(CALLBACK, callback_proc)
+        end
       end
     end
   end

lib/new_relic/supportability_helper.rb

@@ -46,6 +46,7 @@ module SupportabilityHelper
       :recording_web_transaction?,
       :require_test_helper,
       :set_error_group_callback,
+      :set_segment_callback,
       :set_sql_obfuscator,
       :set_transaction_name,
       :set_user_id,

test/multiverse/suites/active_record_pg/before_suite.rb

@@ -26,9 +26,9 @@ def redefine_mysql_primary_key(const_str)
 class Minitest::Test
   def after_teardown
     super
-    User.delete_all
-    Alias.delete_all
-    Order.delete_all
-    Shipment.delete_all
+    User.delete_all if defined?(User)
+    Alias.delete_all if defined?(Alias)
+    Order.delete_all if defined?(Order)
+    Shipment.delete_all if defined?(Shipment)
   end
 end

test/multiverse/suites/active_record_pg/external_request_from_within_ar_block_test.rb

@@ -0,0 +1,57 @@
+# 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 'net/http'
+
+class ExternalRequestFromWithinARBlockTest < Minitest::Test
+  # Use the agent's segment callback system to register a callback for the
+  # ExternalRequestSegment class. Every time that class is initialized, the
+  # callback will be called and it will check to see if the external request
+  # segment has been created from within an ActiveRecord transaction block.
+  # If that check succeeds, generate an error and have the agent notice it.
+  #
+  # https://github.com/newrelic/newrelic-ruby-agent/issues/1556
+  def test_callback_to_notice_error_if_an_external_request_is_made_within_an_ar_block
+    callback = proc do
+      return unless ActiveRecord::Base.connection.transaction_open?
+
+      caller = respond_to?(:name) ? name : '(unknown)'
+      klass = respond_to?(:class) ? self.class.name : '(unknown)'
+      method = __method__ || '(unknown)'
+
+      msg = 'External request made from within an ActiveRecord transaction:' +
+        "\ncaller=#{caller}\nclass=#{klass}\nmethod=#{method}"
+      error = StandardError.new(msg)
+      NewRelic::Agent.notice_error(error)
+    end
+
+    NewRelic::Agent::Transaction::ExternalRequestSegment.set_segment_callback(callback)
+
+    in_transaction do |txn|
+      ActiveRecord::Base.transaction do
+        perform_net_request
+      end
+      external_segments = txn.segments.select { |s| s.name.start_with?('External/') }
+
+      assert_equal 1, external_segments.size
+      segment = external_segments.first
+
+      assert segment, "Failed to find an 'External/' request segment"
+      error = segment.noticed_error
+
+      assert error, "The 'External/' request segment did not contain a noticed error"
+      assert_match 'External request made from within an ActiveRecord transaction', error.message
+    end
+  end
+
+  private
+
+  def perform_net_request
+    uri = URI('https://newrelic.com')
+    http = Net::HTTP.new(uri.host, uri.port)
+    http.use_ssl = true
+    http.verify_mode = OpenSSL::SSL::VERIFY_NONE
+    http.get('/')
+  end
+end

test/new_relic/agent/transaction/abstract_segment_test.rb

@@ -29,6 +29,8 @@ def basic_segment
 
         def setup
           nr_freeze_process_time
+          var = NewRelic::Agent::Transaction::AbstractSegment::CALLBACK
+          BasicSegment.remove_instance_variable(var) if BasicSegment.instance_variable_defined?(var)
         end
 
         def teardown
@@ -318,6 +320,52 @@ def test_range_overlap_for_non_intersecting_ranges
 
           assert_in_delta(0.0, segment.send(:range_overlap, 4.0..5.0))
         end
+
+        # BEGIN callbacks
+        def test_self_set_segment_callback
+          callback = proc { puts 'Hello, World!' }
+          BasicSegment.set_segment_callback(callback)
+
+          assert_equal callback, BasicSegment.instance_variable_get(BasicSegment::CALLBACK)
+        end
+
+        def test_self_set_segment_callback_with_a_non_proc_object
+          skip_unless_minitest5_or_above
+
+          logger = Minitest::Mock.new
+          logger.expect :error, nil, [/expected an argument of type Proc/]
+          NewRelic::Agent.stub :logger, logger do
+            BasicSegment.set_segment_callback([])
+
+            refute BasicSegment.instance_variable_defined?(NewRelic::Agent::Transaction::AbstractSegment::CALLBACK)
+          end
+          logger.verify
+        end
+
+        def test_callback_invocation
+          output = 'Hello, World!'
+          callback = proc { puts output }
+          BasicSegment.set_segment_callback(callback)
+
+          assert_output "#{output}\n" do
+            basic_segment # this calls BasicSegment.new
+          end
+        end
+
+        def test_callback_usage_generated_supportability_metrics
+          skip_unless_minitest5_or_above
+
+          metric = NewRelic::SupportabilityHelper::API_SUPPORTABILITY_METRICS[:set_segment_callback]
+          engine_mock = Minitest::Mock.new
+          engine_mock.expect :tl_record_unscoped_metrics, nil, [metric]
+          NewRelic::Agent.instance.stub :stats_engine, engine_mock do
+            BasicSegment.set_segment_callback(-> { Hash[*%w[hello world]] })
+            basic_segment
+          end
+
+          engine_mock.verify
+        end
+        # END callbacks
       end
     end
   end