diff --git a/lib/ddtrace/contrib/active_support/notifications.rb b/lib/ddtrace/contrib/active_support/notifications.rb index 991cc41..56d9d08 100644 --- a/lib/ddtrace/contrib/active_support/notifications.rb +++ b/lib/ddtrace/contrib/active_support/notifications.rb @@ -1,32 +1,39 @@ module Datadog module Contrib module ActiveSupport - class Notifications - def self.subscribe(pattern, span_name, options = {}, tracer = Datadog.tracer, &block) - subscriber = Subscriber.new(tracer, span_name, options, &block) + # For subscribing and tracing ActiveRecord::Notifications + module Notifications + module_function + def subscribe(pattern, subscriber) ::ActiveSupport::Notifications.subscribe(pattern, subscriber) end + def subscriber(span_name, options = {}, tracer = Datadog.tracer, &block) + Subscriber.new(span_name, options, tracer, &block) + end + + def subscriber!(pattern, span_name, options = {}, tracer = Datadog.tracer, &block) + subscriber(span_name, options, tracer, &block).tap do |subscriber| + subscribe(pattern, subscriber) + end + end + + # Represents a single subscriber/tracer for an ActiveSupport::Notification class Subscriber - def initialize(tracer, span_name, options, &block) - @tracer = tracer + def initialize(span_name, options, tracer, &block) @span_name = span_name @options = options + @tracer = tracer @block = block - - # A stack of open spans. We want to get to the most recently opened span - # at any point in time. - @spans = [] end def start(_name, _id, _payload) - @spans << @tracer.trace(@span_name, @options) + @tracer.trace(@span_name, @options) end def finish(name, id, payload) - # We close spans in reverse order. - span = @spans.pop or raise "no spans left in stack!" + span = @tracer.active_span # The subscriber block needs to remember to set the name of the span. @block.call(span, name, id, payload) diff --git a/lib/ddtrace/contrib/racecar/patcher.rb b/lib/ddtrace/contrib/racecar/patcher.rb index ec95d28..ab8c8c3 100644 --- a/lib/ddtrace/contrib/racecar/patcher.rb +++ b/lib/ddtrace/contrib/racecar/patcher.rb @@ -18,10 +18,7 @@ module Datadog return patched? if patched? || !compatible? tracer = configuration[:tracer] - - Contrib::ActiveSupport::Notifications.subscribe('process_message.racecar', NAME_MESSAGE, {}, tracer, &method(:process)) - Contrib::ActiveSupport::Notifications.subscribe('process_batch.racecar', NAME_BATCH, {}, tracer, &method(:process)) - + subscribe! tracer.set_service_info( configuration[:service_name], 'racecar', @@ -31,12 +28,33 @@ module Datadog @patched = true end - def patched? - return @patched if defined?(@patched) - @patched = false + def subscribe! + return @subscribed if @subscribed + subscribers.each do |pattern, subscriber| + Contrib::ActiveSupport::Notifications.subscribe(pattern, subscriber) + end + @subscribed = true end - def unpatch + def subscribers + @subscribers ||= { + 'process_message.racecar' => Contrib::ActiveSupport::Notifications.subscriber( + self::NAME_MESSAGE, + {}, + configuration[:tracer], + &method(:process) + ), + 'process_batch.racecar' => Contrib::ActiveSupport::Notifications.subscriber( + self::NAME_BATCH, + {}, + configuration[:tracer], + &method(:process) + ) + } + end + + def patched? + return @patched if defined?(@patched) @patched = false end diff --git a/spec/ddtrace/contrib/racecar/patcher_spec.rb b/spec/ddtrace/contrib/racecar/patcher_spec.rb index 368dae7..4e98d7c 100644 --- a/spec/ddtrace/contrib/racecar/patcher_spec.rb +++ b/spec/ddtrace/contrib/racecar/patcher_spec.rb @@ -17,10 +17,14 @@ RSpec.describe 'Racecar patcher' do Datadog.configure do |c| c.use :racecar, tracer: tracer end - end - after(:each) do - Datadog::Contrib::Racecar::Patcher.unpatch + # Make sure to update the subscription tracer, + # so we aren't writing to a stale tracer. + if Datadog::Contrib::Racecar::Patcher.patched? + Datadog::Contrib::Racecar::Patcher.subscribers.each do |_pattern, subscriber| + subscriber.instance_variable_set(:@tracer, tracer) + end + end end describe 'for single message processing' do