Populate ActiveSupport Notification spans as early as possible

lib/datadog/tracing/contrib/action_cable/event.rb

@@ -46,7 +46,7 @@ def self.included(base)
           module ClassMethods
             include Contrib::ActionCable::Event::ClassMethods
 
-            def subscription(*args)
+            def subscription(*args, **kwargs)
               super.tap do |subscription|
                 subscription.before_trace { ensure_clean_context! }
               end

lib/datadog/tracing/contrib/action_cable/events/broadcast.rb

@@ -34,7 +34,7 @@ def span_type
               Tracing::Metadata::Ext::AppTypes::TYPE_WEB
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               channel = payload[:broadcasting] # Channel has high cardinality
               span.service = configuration[:service_name] if configuration[:service_name]
               span.type = span_type

lib/datadog/tracing/contrib/action_cable/events/perform_action.rb

@@ -33,7 +33,7 @@ def span_type
               Tracing::Metadata::Ext::AppTypes::TYPE_WEB
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               channel_class = payload[:channel_class]
               action = payload[:action]
 

lib/datadog/tracing/contrib/action_cable/events/transmit.rb

@@ -33,7 +33,7 @@ def span_type
               Tracing::Metadata::Ext::AppTypes::TYPE_WEB
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               channel_class = payload[:channel_class]
 
               span.service = configuration[:service_name] if configuration[:service_name]

lib/datadog/tracing/contrib/action_mailer/event.rb

@@ -27,9 +27,11 @@ def configuration
               Datadog.configuration.tracing[:action_mailer]
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
+              super
+
+              span.type = span_type
               span.service = configuration[:service_name] if configuration[:service_name]
-              span.resource = payload[:mailer]
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)
 
               # Set analytics sample rate
@@ -39,10 +41,6 @@ def process(span, event, _id, payload)
 
               # Measure service stats
               Contrib::Analytics.set_measured(span)
-
-              report_if_exception(span, payload)
-            rescue StandardError => e
-              Datadog.logger.debug(e.message)
             end
           end
         end

lib/datadog/tracing/contrib/action_mailer/events/deliver.rb

@@ -30,15 +30,20 @@ def span_type
               Tracing::Metadata::Ext::AppTypes::TYPE_WORKER
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               super
 
-              span.type = span_type
+              span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_DELIVER)
+            end
+
+            def on_finish(span, event, _id, payload)
+              super
+
+              span.resource = payload[:mailer] # Mailer is not available at `on_start`
+
               span.set_tag(Ext::TAG_MAILER, payload[:mailer])
               span.set_tag(Ext::TAG_MSG_ID, payload[:message_id])
 
-              span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_DELIVER)
-
               # Since email data can contain PII we disable by default
               # Some of these fields can be either strings or arrays, so we try to normalize
               # https://github.com/rails/rails/blob/18707ab17fa492eb25ad2e8f9818a320dc20b823/actionmailer/lib/action_mailer/base.rb#L742-L754

lib/datadog/tracing/contrib/action_mailer/events/process.rb

@@ -30,10 +30,11 @@ def span_type
               Tracing::Metadata::Ext::HTTP::TYPE_TEMPLATE
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               super
 
-              span.type = span_type
+              span.resource = payload[:mailer] # Mailer is not available at `on_start`
+
               span.set_tag(Ext::TAG_ACTION, payload[:action])
               span.set_tag(Ext::TAG_MAILER, payload[:mailer])
 

lib/datadog/tracing/contrib/action_view/events/render_partial.rb

@@ -27,7 +27,7 @@ def span_name
               Ext::SPAN_RENDER_PARTIAL
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               span.service = configuration[:service_name] if configuration[:service_name]
               span.type = Tracing::Metadata::Ext::HTTP::TYPE_TEMPLATE
 
@@ -41,10 +41,6 @@ def process(span, _event, _id, payload)
 
               # Measure service stats
               Contrib::Analytics.set_measured(span)
-
-              record_exception(span, payload)
-            rescue StandardError => e
-              Datadog.logger.debug(e.message)
             end
           end
         end

lib/datadog/tracing/contrib/action_view/events/render_template.rb

@@ -27,7 +27,7 @@ def span_name
               Ext::SPAN_RENDER_TEMPLATE
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               span.service = configuration[:service_name] if configuration[:service_name]
               span.type = Tracing::Metadata::Ext::HTTP::TYPE_TEMPLATE
 

lib/datadog/tracing/contrib/active_job/events/discard.rb

@@ -26,7 +26,7 @@ def span_name
               Ext::SPAN_DISCARD
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.name = span_name
               span.service = configuration[:service_name] if configuration[:service_name]
               span.resource = payload[:job].class.name

lib/datadog/tracing/contrib/active_job/events/enqueue.rb

@@ -26,7 +26,7 @@ def span_name
               Ext::SPAN_ENQUEUE
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.name = span_name
               span.service = configuration[:service_name] if configuration[:service_name]
               span.resource = payload[:job].class.name

lib/datadog/tracing/contrib/active_job/events/enqueue_at.rb

@@ -26,7 +26,7 @@ def span_name
               Ext::SPAN_ENQUEUE
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.name = span_name
               span.service = configuration[:service_name] if configuration[:service_name]
               span.resource = payload[:job].class.name

lib/datadog/tracing/contrib/active_job/events/enqueue_retry.rb

@@ -26,7 +26,7 @@ def span_name
               Ext::SPAN_ENQUEUE_RETRY
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.name = span_name
               span.service = configuration[:service_name] if configuration[:service_name]
               span.resource = payload[:job].class.name

lib/datadog/tracing/contrib/active_job/events/perform.rb

@@ -26,7 +26,7 @@ def span_name
               Ext::SPAN_PERFORM
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.name = span_name
               span.service = configuration[:service_name] if configuration[:service_name]
               span.resource = payload[:job].class.name

lib/datadog/tracing/contrib/active_job/events/retry_stopped.rb

@@ -26,7 +26,7 @@ def span_name
               Ext::SPAN_RETRY_STOPPED
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.name = span_name
               span.service = configuration[:service_name] if configuration[:service_name]
               span.resource = payload[:job].class.name

lib/datadog/tracing/contrib/active_model_serializers/events/render.rb

@@ -30,7 +30,7 @@ def span_name
               Ext::SPAN_RENDER
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_RENDER)
 
               set_common_tags(span, payload)

lib/datadog/tracing/contrib/active_model_serializers/events/serialize.rb

@@ -32,7 +32,7 @@ def span_name
               Ext::SPAN_SERIALIZE
             end
 
-            def process(span, _event, _id, payload)
+            def on_start(span, _event, _id, payload)
               span.set_tag(Tracing::Metadata::Ext::TAG_OPERATION, Ext::TAG_OPERATION_SERIALIZE)
 
               set_common_tags(span, payload)

lib/datadog/tracing/contrib/active_record/events/instantiation.rb

@@ -31,7 +31,7 @@ def span_name
               Ext::SPAN_INSTANTIATION
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               span.resource = payload.fetch(:class_name)
               span.type = Ext::SPAN_TYPE_INSTANTIATION
               span.set_tag(Tracing::Metadata::Ext::TAG_COMPONENT, Ext::TAG_COMPONENT)

lib/datadog/tracing/contrib/active_record/events/sql.rb

@@ -29,7 +29,7 @@ def span_name
               Ext::SPAN_SQL
             end
 
-            def process(span, event, _id, payload)
+            def on_start(span, event, _id, payload)
               config = Utils.connection_config(payload[:connection], payload[:connection_id])
               settings = Datadog.configuration.tracing[:active_record, config]
               adapter_name = Contrib::Utils::Database.normalize_vendor(config[:adapter])

lib/datadog/tracing/contrib/active_support/notifications/event.rb

@@ -26,21 +26,23 @@ def subscribe! # rubocop:disable Lint/UselessMethodDefinition
                 super
               end
 
-              def subscription(span_name = nil, options = nil)
+              def subscription(span_name = nil, span_options = nil, on_start: nil, on_finish: nil)
                 super(
                   span_name || self.span_name,
-                  options || span_options,
-                  &method(:process)
+                  span_options || self.span_options,
+                  on_start: on_start,
+                  on_finish: on_finish
                 )
               end
 
-              def subscribe(pattern = nil, span_name = nil, options = nil)
+              def subscribe(pattern = nil, span_name = nil, span_options = nil)
                 if supported?
                   super(
                     pattern || event_name,
                     span_name || self.span_name,
-                    options || span_options,
-                    &method(:process)
+                    span_options || self.span_options,
+                    on_start: method(:on_start),
+                    on_finish: method(:on_finish)
                   )
                 end
               end
@@ -62,6 +64,21 @@ def payload_exception(payload)
                 payload[:exception_object] ||
                   payload[:exception] # Fallback for ActiveSupport < 5.0
               end
+
+              def on_start(_span, _event, _id, _payload); end
+
+              def on_finish(span, _event, _id, payload)
+                record_exception(span, payload)
+              end
+
+              def record_exception(span, payload)
+                if payload[:exception_object]
+                  span.set_error(payload[:exception_object])
+                elsif payload[:exception]
+                  # Fallback for ActiveSupport < 5.0
+                  span.set_error(payload[:exception])
+                end
+              end
             end
           end
         end

lib/datadog/tracing/contrib/active_support/notifications/subscriber.rb

@@ -45,16 +45,16 @@ def subscribe!
               end
 
               # Creates a subscription and immediately activates it.
-              def subscribe(pattern, span_name, options = {}, &block)
-                subscription(span_name, options, &block).tap do |subscription|
+              def subscribe(pattern, span_name, span_options = {}, on_start: nil, on_finish: nil)
+                subscription(span_name, span_options, on_start: on_start, on_finish: on_finish).tap do |subscription|
                   subscription.subscribe(pattern)
                 end
               end
 
               # Creates a subscription without activating it.
               # Subscription is added to the inheriting class' list of subscriptions.
-              def subscription(span_name, options = {}, &block)
-                Subscription.new(span_name, options, &block).tap do |subscription|
+              def subscription(span_name, span_options = {}, on_start: nil, on_finish: nil)
+                Subscription.new(span_name, span_options, on_start: on_start, on_finish: on_finish).tap do |subscription|
                   subscriptions << subscription
                 end
               end

lib/datadog/tracing/contrib/active_support/notifications/subscription.rb

@@ -9,14 +9,21 @@ module Notifications
           class Subscription
             attr_accessor \
               :span_name,
-              :options
+              :span_options
 
-            def initialize(span_name, options, &block)
-              raise ArgumentError, 'Must be given a block!' unless block
+            # @param span_name [String] the operation name for the span
+            # @param span_options [Hash] span_options to pass during span creation
+            # @param on_start [Proc] a block to run when the event is fired,
+            #   might not include all required information in the `payload` argument.
+            # @param on_finish [Proc] a block to run when the event has finished processing,
+            #   possibly including more information in the `payload` argument.
+            def initialize(span_name, span_options, on_start: nil, on_finish: nil)
+              raise ArgumentError, 'Must be given either on_start or on_finish' unless on_start || on_finish
 
               @span_name = span_name
-              @options = options
-              @handler = Handler.new(&block)
+              @span_options = span_options
+              @on_start = Handler.new(on_start)
+              @on_finish = Handler.new(on_finish)
               @callbacks = Callbacks.new
             end
 
@@ -69,19 +76,24 @@ def unsubscribe_all
             protected
 
             attr_reader \
-              :handler,
+              :on_start,
+              :on_finish,
               :callbacks
 
             def start_span(name, id, payload, start = nil)
               # Run callbacks
               callbacks.run(name, :before_trace, id, payload, start)
 
               # Start a trace
-              Tracing.trace(@span_name, **@options).tap do |span|
-                # Start span if time is provided
-                span.start(start) unless start.nil?
-                payload[:datadog_span] = span
-              end
+              span = Tracing.trace(@span_name, **@span_options)
+
+              # Start span if time is provided
+              span.start(start) unless start.nil?
+              payload[:datadog_span] = span
+
+              on_start.run(span, name, id, payload)
+
+              span
             end
 
             def finish_span(name, id, payload, finish = nil)
@@ -90,7 +102,7 @@ def finish_span(name, id, payload, finish = nil)
                 return nil if span.nil?
 
                 # Run handler for event
-                handler.run(span, name, id, payload)
+                on_finish.run(span, name, id, payload)
 
                 # Finish the span
                 span.finish(finish)
@@ -109,21 +121,17 @@ def subscribers
             class Handler
               attr_reader :block
 
-              def initialize(&block)
+              def initialize(block)
                 @block = block
               end
 
               def run(span, name, id, payload)
-                run!(span, name, id, payload)
+                @block.call(span, name, id, payload) if @block
               rescue StandardError => e
                 Datadog.logger.debug(
                   "ActiveSupport::Notifications handler for '#{name}' failed: #{e.class.name} #{e.message}"
                 )
               end
-
-              def run!(*args)
-                @block.call(*args)
-              end
             end
 
             # Wrapper for subscription callbacks

lib/datadog/tracing/contrib/kafka/consumer_event.rb

@@ -6,7 +6,7 @@ module Contrib
       module Kafka
         # Defines basic behaviors for an event for a consumer.
         module ConsumerEvent
-          def process(span, _event, _id, payload)
+          def on_start(span, _event, _id, payload)
             super
 
             span.set_tag(Ext::TAG_GROUP, payload[:group_id])