Extract context clearing logic for ActiveSupport events

lib/ddtrace/contrib/action_cable/event.rb

@@ -27,39 +27,6 @@ def configuration
           end
         end
       end
-
-      # Defines behavior for the first event of a thread execution.
-      #
-      # This event is not expected to be nested with other event,
-      # but to start a fresh tracing context.
-      module RootContextEvent
-        def self.included(base)
-          base.send(:include, ActiveSupport::Notifications::Event)
-          base.send(:extend, ClassMethods)
-        end
-
-        # Class methods for events.
-        module ClassMethods
-          include Contrib::ActionCable::Event::ClassMethods
-
-          def subscription(*args)
-            super.tap do |subscription|
-              subscription.before_trace { ensure_clean_context! }
-            end
-          end
-
-          private
-
-          # Context objects are thread-bound.
-          # If an integration re-uses threads, context from a previous trace
-          # could leak into the new trace. This "cleans" current context,
-          # preventing such a leak.
-          def ensure_clean_context!
-            return unless configuration[:tracer].call_context.current_span
-            configuration[:tracer].provider.context = Context.new
-          end
-        end
-      end
     end
   end
 end

lib/ddtrace/contrib/action_cable/events/perform_action.rb

@@ -10,7 +10,8 @@ module Events
         # An action, triggered by a WebSockets client, invokes a method
         # in the server's channel instance.
         module PerformAction
-          include ActionCable::RootContextEvent
+          include ActionCable::Event
+          extend ActiveSupport::Notifications::RootEvent
 
           EVENT_NAME = 'perform_action.action_cable'.freeze
 

lib/ddtrace/contrib/active_support/notifications/event.rb

@@ -56,6 +56,40 @@ def tracer
             end
           end
         end
+
+        # Extension to {Event} class that ensures the current {Context}
+        # is always clean when the event is processed.
+        #
+        # This is a safeguard as Contexts are thread-bound.
+        # If an integration re-uses threads, the context from a previous
+        # execution could leak into the new execution.
+        #
+        # This module *cannot* be used for events can be nested, as
+        # it drops all spans currently active in the {Context}.
+        module RootEvent
+          def subscription(*args)
+            super.tap do |subscription|
+              subscription.before_trace { ensure_clean_context! }
+            end
+          end
+
+          private
+
+          # Clears context if there are unfinished spans in it
+          def ensure_clean_context!
+            unfinished_span = configuration[:tracer].call_context.current_span
+            return unless unfinished_span
+
+            Diagnostics::Health.metrics.error_unfinished_context(
+              1, tags: [
+                "span_name:#{unfinished_span.name}",
+                "event:#{self}"
+              ]
+            )
+
+            configuration[:tracer].provider.context = Context.new
+          end
+        end
       end
     end
   end

lib/ddtrace/contrib/racecar/event.rb

@@ -10,17 +10,12 @@ module Event
         def self.included(base)
           base.send(:include, ActiveSupport::Notifications::Event)
           base.send(:extend, ClassMethods)
+          base.send(:extend, ActiveSupport::Notifications::RootEvent)
         end
 
         # Class methods for Racecar events.
         # Note, they share the same process method and before_trace method.
         module ClassMethods
-          def subscription(*args)
-            super.tap do |subscription|
-              subscription.before_trace { ensure_clean_context! }
-            end
-          end
-
           def span_options
             { service: configuration[:service_name] }
           end
@@ -49,17 +44,6 @@ def process(span, event, _id, payload)
             span.set_tag(Ext::TAG_MESSAGE_COUNT, payload[:message_count]) if payload.key?(:message_count)
             span.set_error(payload[:exception_object]) if payload[:exception_object]
           end
-
-          private
-
-          # Context objects are thread-bound.
-          # If Racecar re-uses threads, context from a previous trace
-          # could leak into the new trace. This "cleans" current context,
-          # preventing such a leak.
-          def ensure_clean_context!
-            return unless configuration[:tracer].call_context.current_span
-            configuration[:tracer].provider.context = Context.new
-          end
         end
       end
     end

lib/ddtrace/diagnostics/health.rb

@@ -13,6 +13,7 @@ class Metrics < ::Datadog::Metrics
         count :error_context_overflow, Ext::Diagnostics::Health::Metrics::METRIC_ERROR_CONTEXT_OVERFLOW
         count :error_instrumentation_patch, Ext::Diagnostics::Health::Metrics::METRIC_ERROR_INSTRUMENTATION_PATCH
         count :error_span_finish, Ext::Diagnostics::Health::Metrics::METRIC_ERROR_SPAN_FINISH
+        count :error_unfinished_context, Ext::Diagnostics::Health::Metrics::METRIC_ERROR_UNFINISHED_CONTEXT
         count :error_unfinished_spans, Ext::Diagnostics::Health::Metrics::METRIC_ERROR_UNFINISHED_SPANS
         count :instrumentation_patched, Ext::Diagnostics::Health::Metrics::METRIC_INSTRUMENTATION_PATCHED
         count :queue_accepted, Ext::Diagnostics::Health::Metrics::METRIC_QUEUE_ACCEPTED

lib/ddtrace/ext/diagnostics.rb

@@ -13,6 +13,7 @@ module Metrics
           METRIC_ERROR_CONTEXT_OVERFLOW = 'datadog.tracer.error.context_overflow'.freeze
           METRIC_ERROR_INSTRUMENTATION_PATCH = 'datadog.tracer.error.instrumentation_patch'.freeze
           METRIC_ERROR_SPAN_FINISH = 'datadog.tracer.error.span_finish'.freeze
+          METRIC_ERROR_UNFINISHED_CONTEXT = 'datadog.tracer.error.unfinished_context'.freeze
           METRIC_ERROR_UNFINISHED_SPANS = 'datadog.tracer.error.unfinished_spans'.freeze
           METRIC_INSTRUMENTATION_PATCHED = 'datadog.tracer.instrumentation_patched'.freeze
           METRIC_QUEUE_ACCEPTED = 'datadog.tracer.queue.accepted'.freeze

spec/ddtrace/contrib/action_cable/patcher_spec.rb

@@ -2,6 +2,7 @@
 
 require 'ddtrace'
 require 'ddtrace/contrib/analytics_examples'
+require 'ddtrace/contrib/tracer_examples'
 
 require 'rails'
 require 'active_support/core_ext/hash/indifferent_access'
@@ -105,8 +106,10 @@ def foo(_data); end
         let(:analytics_sample_rate_var) { Datadog::Contrib::ActionCable::Ext::ENV_ANALYTICS_SAMPLE_RATE }
       end
 
-      context 'with a leaking context' do
-        let!(:leaky_span) { tracer.trace('unfinished_span') }
+      context 'with an unfinished trace' do
+        include_context 'an unfinished trace'
+
+        let(:event_name) { 'Datadog::Contrib::ActionCable::Events::PerformAction' }
 
         it 'traces transmit event' do
           perform

spec/ddtrace/contrib/racecar/patcher_spec.rb

@@ -1,5 +1,6 @@
 require 'spec_helper'
 require 'ddtrace/contrib/analytics_examples'
+require 'ddtrace/contrib/tracer_examples'
 
 require 'racecar'
 require 'racecar/cli'
@@ -91,6 +92,18 @@ def all_spans
       end
     end
 
+    context 'with an unfinished trace' do
+      include_context 'an unfinished trace'
+
+      let(:event_name) { 'Datadog::Contrib::Racecar::Events::Message' }
+
+      it 'traces transmit event' do
+        ActiveSupport::Notifications.instrument('process_message.racecar', payload)
+
+        expect(span.name).to eq('racecar.message')
+      end
+    end
+
     it_behaves_like 'analytics for integration' do
       before { ActiveSupport::Notifications.instrument('process_message.racecar', payload) }
       let(:analytics_enabled_var) { Datadog::Contrib::Racecar::Ext::ENV_ANALYTICS_ENABLED }

spec/ddtrace/contrib/tracer_examples.rb

@@ -0,0 +1,12 @@
+RSpec.shared_context 'an unfinished trace' do
+  let(:unfinished_span_name) { 'unfinished_span' }
+  let!(:unfinished_span) { tracer.trace(unfinished_span_name) }
+
+  before do
+    expect(Datadog::Diagnostics::Health.metrics).to receive(:error_unfinished_context)
+      .with(1, tags: [
+              "span_name:#{unfinished_span_name}",
+              "event:#{event_name}"
+            ])
+  end
+end