instrumentation.rb in newrelic_rpm-8.11.0

- old
+ new

@@ -1,17 +1,15 @@
-# encoding: utf-8
 # 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
   module Agent
     module Instrumentation
       module Curb
         module Easy
           attr_accessor :_nr_instrumented,
-            :_nr_failure_instrumented,
             :_nr_http_verb,
             :_nr_header_str,
             :_nr_original_on_header,
             :_nr_original_on_complete,
             :_nr_original_on_failure,
@@ -141,11 +139,11 @@
           def install_completion_callback(request, wrapped_response, segment)
             original_callback = request.on_complete
             request._nr_original_on_complete = original_callback
             request.on_complete do |finished_request|
               begin
-                segment.process_response_headers(wrapped_response)
+                segment.process_response_headers(wrapped_response) if segment
               ensure
                 segment.finish if segment
                 # Make sure the existing completion callback is run, and restore the
                 # on_complete callback to how it was before.
                 original_callback.call(finished_request) if original_callback
@@ -156,28 +154,34 @@
 
           # Install a callback that will fire on failures
           # NOTE:  on_failure is not always called, so we're not always
           # unhooking the callback.  No harm/no foul in production, but
           # definitely something to beware of if debugging callback issues
-          # _nr_failure_instrumented exists to prevent infinitely chaining
+          # @__newrelic_original_callback exists to prevent infinitely chaining
           # our on_failure callback hook.
-          def install_failure_callback(request, wrapped_response, segment)
-            return if request._nr_failure_instrumented
+          def install_failure_callback(request, _wrapped_response, segment)
             original_callback = request.on_failure
+
+            nr_original_callback = original_callback.instance_variable_get(:@__newrelic_original_callback)
+            original_callback = nr_original_callback || original_callback
+
             request._nr_original_on_failure = original_callback
-            request.on_failure do |failed_request, error|
+
+            newrelic_callback = proc do |failed_request, error|
               begin
                 if segment
-                  noticible_error = NewRelic::Agent::NoticibleError.new(error[0].name, error[-1])
-                  segment.notice_error(noticible_error)
+                  noticeable_error = NewRelic::Agent::NoticeableError.new(error[0].name, error[-1])
+                  segment.notice_error(noticeable_error)
                 end
               ensure
                 original_callback.call(failed_request, error) if original_callback
                 remove_failure_callback(failed_request)
               end
-              request._nr_failure_instrumented = true
             end
+            newrelic_callback.instance_variable_set(:@__newrelic_original_callback, original_callback)
+
+            request.on_failure(&newrelic_callback)
           end
 
           # on_failure callbacks cannot be removed in the on_complete
           # callback where this method is invoked because on_complete
           # fires before the on_failure!
@@ -190,10 +194,9 @@
           # We execute customer's on_failure callback (if any) and
           # uninstall our hook here since the on_complete callback
           # fires before the on_failure callback.
           def remove_failure_callback(request)
             request.on_failure(&request._nr_original_on_failure)
-            request._nr_failure_instrumented = false
           end
 
           private
 
           def first_request_is_serial?