Add automatic instrumentation support for callbacks

bullet_train/lib/bullet_train/core_ext/callbacks_instrumentation.rb

@@ -0,0 +1,60 @@
+require_relative "callbacks_instrumentation/backport"
+
+module ActiveSupport::Callbacks
+ ClassMethods.prepend Module.new {
+ def set_callback(name, *filter_list, &block)
+ *, options = normalize_callback_params(filter_list.dup, block)
+ Instrumenter.options = options.merge(caller_locations:)
+ super
+ ensure
+ Instrumenter.options = nil
+ end
+ }
+
+ Callback.prepend Module.new {
+ def compiled
+ @compiled = super.then { Instrumenter.build self, _1 } unless defined?(@compiled)
+ @compiled
+ end
+ }
+
+ class Instrumenter < Data.define(:callback, :callable, :options)
+ singleton_class.attr_accessor :options
+
+ def self.build(callback, callable)
+ options ? new(callback, callable, options) : callable
+ end
+
+ def initialize(callback:, callable:, options:)
+ caller_locations = options.delete(:caller_locations)
+ @location = Rails.backtrace_cleaner.clean(caller_locations.map(&:to_s)).first if caller_locations
+ @full_name = (callback.name == :commit) ? caller_locations.first.base_label : "#{callback.kind}_#{callback.name}"
+ super
+ end
+ attr_reader :options, :location, :full_name
+ delegate :name, :kind, :filter, to: :callback
+
+ def apply(sequence)
+ case kind
+ in :before then sequence.before(self)
+ in :after then sequence.after(self)
+ in :around then sequence.around(self)
+ end
+ end
+
+ def call(env)
+ ActiveSupport::Notifications.instrument "callback.active_support", callback: self do
+ callable.call(env)
+ end
+ end
+ end
+
+ class LogSubscriber < ActiveSupport::LogSubscriber
+ attach_to :active_support
+
+ def callback(event)
+ event.payload => {callback:}
+ info { "Callback #{callback.full_name}: #{callback.filter} #{callback.options} from #{callback.location}".squish }
+ end
+ end
+end

bullet_train/lib/bullet_train/core_ext/callbacks_instrumentation/backport.rb

@@ -0,0 +1,118 @@
+return if ActiveSupport::VERSION.to_s.start_with?("7.2")
+
+# Backports a few commits from Rails 7.2 to make our debugging extension easier:
+#
+# - https://github.com/rails/rails/commit/6c5a042824fdaf630ebaa7fb293b8c2543f20d00
+# - https://github.com/rails/rails/commit/42ad4d6b0bfc14e16050f9945a8b2ac5c87c294f
+module ActiveSupport::Callbacks
+ Callback.prepend Module.new {
+ def initialize(...)
+ super
+ compiled # Eager load ActiveSupport::Callback procs
+ end
+
+ def compiled
+ @compiled ||=
+ begin
+ user_conditions = conditions_lambdas
+ user_callback = CallTemplate.build(@filter, self)
+
+ case kind
+ when :before
+ Filters::Before.new(user_callback.make_lambda, user_conditions, chain_config, @filter, name)
+ when :after
+ Filters::After.new(user_callback.make_lambda, user_conditions, chain_config)
+ when :around
+ Filters::Around.new(user_callback, user_conditions)
+ end
+ end
+ end
+
+ def apply(callback_sequence)
+ compiled.apply(callback_sequence)
+ end
+ }
+
+ class CallbackSequence
+ def before(before)
+ @before.unshift(before)
+ self
+ end
+
+ def after(after)
+ @after.push(after)
+ self
+ end
+ end
+
+ module Filters
+ # Around wasn't defined before our commit backport, so we don't need to remove that.
+ remove_const :Before
+ remove_const :After
+
+ class Before
+ def initialize(user_callback, user_conditions, chain_config, filter, name)
+ halted_lambda = chain_config[:terminator]
+ @user_callback, @user_conditions, @halted_lambda, @filter, @name = user_callback, user_conditions, halted_lambda, filter, name
+ freeze
+ end
+ attr_reader :user_callback, :user_conditions, :halted_lambda, :filter, :name
+
+ def call(env)
+ target = env.target
+ value = env.value
+ halted = env.halted
+
+ if !halted && user_conditions.all? { |c| c.call(target, value) }
+ result_lambda = -> { user_callback.call target, value }
+ env.halted = halted_lambda.call(target, result_lambda)
+ if env.halted
+ target.send :halted_callback_hook, filter, name
+ end
+ end
+
+ env
+ end
+
+ def apply(callback_sequence)
+ callback_sequence.before(self)
+ end
+ end
+
+ class After
+ attr_reader :user_callback, :user_conditions, :halting
+ def initialize(user_callback, user_conditions, chain_config)
+ halting = chain_config[:skip_after_callbacks_if_terminated]
+ @user_callback, @user_conditions, @halting = user_callback, user_conditions, halting
+ freeze
+ end
+
+ def call(env)
+ target = env.target
+ value = env.value
+ halted = env.halted
+
+ if (!halted || !@halting) && user_conditions.all? { |c| c.call(target, value) }
+ user_callback.call target, value
+ end
+
+ env
+ end
+
+ def apply(callback_sequence)
+ callback_sequence.after(self)
+ end
+ end
+
+ class Around
+ def initialize(user_callback, user_conditions)
+ @user_callback, @user_conditions = user_callback, user_conditions
+ freeze
+ end
+
+ def apply(callback_sequence)
+ callback_sequence.around(@user_callback, @user_conditions)
+ end
+ end
+ end
+end

bullet_train/lib/bullet_train/engine.rb

@@ -21,5 +21,15 @@
 
 module BulletTrain
  class Engine < ::Rails::Engine
+ config.bullet_train = ActiveSupport::OrderedOptions.new
+ config.bullet_train.instrument_callbacks = Rails.env.development? || Rails.env.test?
+
+ initializer "callbacks.instrumentation" do
+ if config.bullet_train.instrument_callbacks
+ config.after_initialize do
+ require_relative "../bullet_train/core_ext/callbacks_instrumentation"
+ end
+ end
+ end
  end
 end