lib/datadog/tracing/contrib/active_support/notifications/event.rb in datadog-2.1.0 vs lib/datadog/tracing/contrib/active_support/notifications/event.rb in datadog-2.2.0

@@ -24,25 +24,29 @@
               # Publicly exposes protected method `subscribe!`
               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, trace: 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,
+                  trace: trace
                 )
               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),
+                    trace: method(:trace?)
                   )
                 end
               end
 
               def supported?
@@ -59,9 +63,28 @@
               end
 
               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 trace?(_event, _payload)
+                true
+              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
       end