Fix trace_id mismatch for logs emitted before Sentry::Rails::CaptureExceptions runs

sentry-rails/lib/sentry/rails/capture_context.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module Sentry
+  module Rails
+    # A minimal middleware that establishes Sentry's trace/scope context as early
+    # as possible in the middleware stack.
+    #
+    # +Sentry::Rails::CaptureExceptions+ is intentionally positioned right after
+    # +ActionDispatch::ShowExceptions+ so it can skip transaction creation for
+    # static asset requests served by earlier middlewares (like +Rack::Sendfile+
+    # and +ActionDispatch::Static+). Because of that placement, anything logged
+    # before it runs - most notably +Rails::Rack::Logger+'s "Started ..." line -
+    # would otherwise get a trace_id/span_id unrelated to the rest of the request.
+    #
+    # This middleware only establishes the propagation context (using the
+    # incoming `sentry-trace`/`baggage` headers, if present); it does not start a
+    # transaction or capture exceptions. +Sentry::Rack::CaptureExceptions+ detects
+    # that context has already been established for this request (via
+    # +PropagationContext::ESTABLISHED_ENV_KEY+) and reuses it instead of
+    # generating a new, mismatched one.
+    class CaptureContext
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        return @app.call(env) unless Sentry.initialized?
+
+        Sentry.clone_hub_to_current_thread
+        Sentry.get_current_scope.generate_propagation_context(env)
+        env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY] = true
+
+        @app.call(env)
+      end
+    end
+  end
+end

sentry-rails/lib/sentry/rails/railtie.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require "sentry/rails/capture_context"
 require "sentry/rails/capture_exceptions"
 require "sentry/rails/rescued_exception_interceptor"
 require "sentry/rails/backtrace_cleaner"
@@ -8,6 +9,10 @@ module Sentry
   class Railtie < ::Rails::Railtie
     # middlewares can't be injected after initialize
     initializer "sentry.use_rack_middleware" do |app|
+      # placed as the very first middleware so that anything logged before CaptureExceptions
+      # runs (e.g. Rails::Rack::Logger's "Started ..." line) shares the same trace context as
+      # the rest of the request
+      app.config.middleware.unshift Sentry::Rails::CaptureContext
       # placed after all the file-sending middlewares so we can avoid unnecessary transactions
       app.config.middleware.insert_after ActionDispatch::ShowExceptions, Sentry::Rails::CaptureExceptions
       # need to place as close to DebugExceptions as possible to intercept most of the exceptions, including those raised by middlewares

sentry-rails/spec/sentry/rails/capture_context_spec.rb

@@ -0,0 +1,103 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe Sentry::Rails::CaptureContext do
+  # Records the current scope's trace_id every time it's called, so specs can
+  # compare what a piece of middleware would see at different points in the stack.
+  class CaptureContextSpecProbe
+    def self.captured_trace_ids
+      @captured_trace_ids ||= []
+    end
+
+    def initialize(app)
+      @app = app
+    end
+
+    def call(env)
+      self.class.captured_trace_ids << Sentry.get_current_scope.get_trace_context[:trace_id]
+      @app.call(env)
+    end
+  end
+
+  describe "#call" do
+    before do
+      make_basic_app
+    end
+
+    it "establishes a propagation context and flags the env" do
+      trace_id_in_app = nil
+
+      app = lambda do |env|
+        trace_id_in_app = Sentry.get_current_scope.get_trace_context[:trace_id]
+        [200, {}, ["ok"]]
+      end
+
+      env = Rack::MockRequest.env_for("/test")
+      described_class.new(app).call(env)
+
+      expect(env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY]).to eq(true)
+      expect(trace_id_in_app).to be_a(String)
+    end
+
+    it "is a no-op when Sentry is not initialized" do
+      allow(Sentry).to receive(:initialized?).and_return(false)
+
+      called = false
+      app = lambda do |env|
+        called = true
+        [200, {}, ["ok"]]
+      end
+
+      env = Rack::MockRequest.env_for("/test")
+      described_class.new(app).call(env)
+
+      expect(called).to eq(true)
+      expect(env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY]).to be_nil
+    end
+  end
+
+  context "when composed with CaptureExceptions", type: :request do
+    before do
+      CaptureContextSpecProbe.captured_trace_ids.clear
+    end
+
+    context "without tracing enabled" do
+      before do
+        make_basic_app do |config, app|
+          app.config.middleware.insert_before(Sentry::Rails::CaptureExceptions, CaptureContextSpecProbe)
+          app.config.middleware.insert_after(Sentry::Rails::CaptureExceptions, CaptureContextSpecProbe)
+        end
+      end
+
+      it "keeps the same trace_id before and after CaptureExceptions runs" do
+        get "/world"
+
+        early_trace_id, late_trace_id = CaptureContextSpecProbe.captured_trace_ids
+
+        expect(early_trace_id).to be_a(String)
+        expect(late_trace_id).to eq(early_trace_id)
+      end
+    end
+
+    context "with tracing enabled" do
+      before do
+        make_basic_app do |config, app|
+          config.traces_sample_rate = 1.0
+          app.config.middleware.insert_before(Sentry::Rails::CaptureExceptions, CaptureContextSpecProbe)
+          app.config.middleware.insert_after(Sentry::Rails::CaptureExceptions, CaptureContextSpecProbe)
+        end
+      end
+
+      it "keeps the same trace_id from before CaptureExceptions through the started transaction" do
+        get "/world"
+
+        early_trace_id, late_trace_id = CaptureContextSpecProbe.captured_trace_ids
+
+        expect(early_trace_id).to be_a(String)
+        # the "late" trace_id comes from the actual transaction CaptureExceptions started
+        expect(late_trace_id).to eq(early_trace_id)
+      end
+    end
+  end
+end

sentry-rails/spec/sentry/rails_spec.rb

@@ -22,6 +22,7 @@
 
     it "inserts middleware to a correct position" do
       app = Rails.application
+      expect(app.middleware.first).to eq(Sentry::Rails::CaptureContext)
       index_of_executor = app.middleware.find_index { |m| m == ActionDispatch::ShowExceptions }
       expect(app.middleware.find_index(Sentry::Rails::CaptureExceptions)).to eq(index_of_executor + 1)
       index_of_debug_exceptions = app.middleware.find_index { |m| m == ActionDispatch::DebugExceptions }

sentry-ruby/lib/sentry/hub.rb

@@ -122,6 +122,22 @@ def start_transaction(transaction: nil, custom_sampling_context: {}, instrumente
       return unless configuration.tracing_enabled?
       return unless instrumenter == configuration.instrumenter
 
+      if transaction.nil? && !options.key?(:trace_id) && trace_context_established?(custom_sampling_context)
+        # An earlier point in this exact request/job (e.g. Sentry::Rails::CaptureContext,
+        # via Sentry::Rack::CaptureExceptions) already established the scope's propagation
+        # context specifically for this operation - most likely because something was
+        # already logged/traced with it before this transaction was created. Adopt its
+        # trace_id (and correspondingly-derived sample_rand) instead of generating an
+        # unrelated one, so this transaction doesn't diverge from what was already
+        # established. This is intentionally scoped to only when the established flag is
+        # present, since the scope always has *some* propagation context by default and
+        # blindly reusing it here would incorrectly link together unrelated transactions
+        # (e.g. separate background jobs sharing a thread).
+        propagation_context = current_scope.propagation_context
+        options[:trace_id] = propagation_context.trace_id
+        options[:sample_rand] ||= propagation_context.sample_rand
+      end
+
       transaction ||= Transaction.new(**options)
 
       sampling_context = {
@@ -374,7 +390,11 @@ def get_trace_propagation_meta
     end
 
     def continue_trace(env, **options)
-      configure_scope { |s| s.generate_propagation_context(env) }
+      # Don't clobber context that an earlier point in the stack already established
+      # for this exact request/env (see PropagationContext::ESTABLISHED_ENV_KEY).
+      unless env && env[PropagationContext::ESTABLISHED_ENV_KEY]
+        configure_scope { |s| s.generate_propagation_context(env) }
+      end
 
       return nil unless configuration.tracing_enabled?
 
@@ -393,6 +413,11 @@ def continue_trace(env, **options)
 
     private
 
+    def trace_context_established?(custom_sampling_context)
+      env = custom_sampling_context[:env]
+      !!(env && env[PropagationContext::ESTABLISHED_ENV_KEY])
+    end
+
     def current_layer
       @stack.last
     end

sentry-ruby/lib/sentry/propagation_context.rb

@@ -13,6 +13,21 @@ class PropagationContext
       "-?([01])?\\z"         # sampled
     )
 
+    # Rack env key used to signal that trace/scope context has already been established
+    # for the current request by an earlier point in the middleware stack (e.g. by
+    # +Sentry::Rails::CaptureContext+, which runs ahead of +Sentry::Rack::CaptureExceptions+
+    # so that anything logged before the latter runs still shares the request's trace_id).
+    #
+    # When this flag is present, +Hub#continue_trace+ and +Sentry::Rack::CaptureExceptions+
+    # reuse the already-established context instead of regenerating/discarding it.
+    #
+    # +Sentry::Rack::CaptureExceptions+ deletes this key from +env+ once it's done using it,
+    # so it only ever affects the single request it was set up for. This matters for
+    # integrations (like Action Cable) that hold on to the same +env+ for the lifetime of a
+    # long-running connection and reuse it across many separate operations - each of those
+    # needs its own freshly-established context rather than perpetually reusing this one.
+    ESTABLISHED_ENV_KEY = "sentry.trace_context_established"
+
     # An uuid that can be used to identify a trace.
     # @return [String]
     attr_reader :trace_id

sentry-ruby/lib/sentry/rack/capture_exceptions.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "sentry/propagation_context"
+
 module Sentry
   module Rack
     class CaptureExceptions
@@ -14,8 +16,10 @@ def initialize(app)
       def call(env)
         return @app.call(env) unless Sentry.initialized?
 
-        # make sure the current thread has a clean hub
-        Sentry.clone_hub_to_current_thread
+        # make sure the current thread has a clean hub, unless an earlier point in the
+        # middleware stack already established one for this exact request (e.g.
+        # Sentry::Rails::CaptureContext) - in that case reuse it instead of discarding it.
+        Sentry.clone_hub_to_current_thread unless env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY]
 
         Sentry.with_scope do |scope|
           Sentry.with_session_tracking do
@@ -26,6 +30,14 @@ def call(env)
             transaction = start_transaction(env, scope)
             scope.set_span(transaction) if transaction
 
+            # Consume the flag now that it's served its purpose (letting continue_trace/
+            # start_transaction above reuse the context established earlier in the stack).
+            # It must not leak into the downstream app call: some integrations (e.g. Action
+            # Cable) hold on to this exact env for the lifetime of a long-running connection
+            # and reuse it across many unrelated operations, each of which needs its own
+            # freshly-established propagation context rather than perpetually reusing this one.
+            env.delete(Sentry::PropagationContext::ESTABLISHED_ENV_KEY)
+
             begin
               response = @app.call(env)
             rescue Sentry::Error

sentry-ruby/spec/sentry/rack/capture_exceptions_spec.rb

@@ -93,6 +93,70 @@
       expect(env.key?("sentry.error_event_id")).to eq(false)
     end
 
+    context "when trace context was already established earlier in the stack" do
+      it "does not re-clone the hub and reuses the existing propagation context" do
+        Sentry.clone_hub_to_current_thread
+        Sentry.get_current_scope.generate_propagation_context(env)
+        env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY] = true
+
+        established_propagation_context = Sentry.get_current_scope.propagation_context
+
+        expect(Sentry).not_to receive(:clone_hub_to_current_thread)
+
+        trace_id_in_app = nil
+        app = lambda do |e|
+          trace_id_in_app = Sentry.get_current_scope.get_trace_context[:trace_id]
+          [200, {}, ['okay']]
+        end
+
+        stack = Sentry::Rack::CaptureExceptions.new(app)
+        stack.call(env)
+
+        expect(trace_id_in_app).to eq(established_propagation_context.trace_id)
+      end
+
+      it "deletes the established flag from env so it doesn't leak into later reuses of the same env" do
+        Sentry.clone_hub_to_current_thread
+        Sentry.get_current_scope.generate_propagation_context(env)
+        env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY] = true
+
+        app = ->(_e) { [200, {}, ['okay']] }
+        stack = Sentry::Rack::CaptureExceptions.new(app)
+        stack.call(env)
+
+        expect(env.key?(Sentry::PropagationContext::ESTABLISHED_ENV_KEY)).to eq(false)
+      end
+
+      it "does not reuse a stale established context on a later, unrelated call with the same env" do
+        # Simulates a long-lived connection (e.g. Action Cable) that stores the handshake's
+        # env and reuses it for many separate operations over its lifetime - only the very
+        # first operation immediately following CaptureContext should honor the flag.
+        Sentry.clone_hub_to_current_thread
+        Sentry.get_current_scope.generate_propagation_context(env)
+        env[Sentry::PropagationContext::ESTABLISHED_ENV_KEY] = true
+
+        app = ->(_e) { [200, {}, ['okay']] }
+        stack = Sentry::Rack::CaptureExceptions.new(app)
+        stack.call(env)
+
+        Sentry.clone_hub_to_current_thread
+        propagation_context_before_second_call = Sentry.get_current_scope.propagation_context
+
+        trace_id_in_second_call = nil
+        second_app = lambda do |e|
+          trace_id_in_second_call = Sentry.get_current_scope.get_trace_context[:trace_id]
+          [200, {}, ['okay']]
+        end
+
+        expect(Sentry).to receive(:clone_hub_to_current_thread).and_call_original
+
+        second_stack = Sentry::Rack::CaptureExceptions.new(second_app)
+        second_stack.call(env)
+
+        expect(trace_id_in_second_call).not_to eq(propagation_context_before_second_call.trace_id)
+      end
+    end
+
     context "with config.include_local_variables = true" do
       before do
         perform_basic_setup do |config|