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

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

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module Sentry
+  module Rails
+    # Establishes the propagation context as early as possible, so anything
+    # logged before +CaptureExceptions+ runs shares the request's trace_id.
+    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/capture_exceptions.rb

@@ -35,7 +35,7 @@ def capture_exception(exception, env)
         end
       end
 
-      def start_transaction(env, scope)
+      def start_transaction(env, scope, established)
         options = {
           name: scope.transaction_name,
           source: scope.transaction_source,
@@ -45,8 +45,8 @@ def start_transaction(env, scope)
 
         options.merge!(sampled: false) if @assets_regexp && scope.transaction_name.match?(@assets_regexp)
 
-        transaction = Sentry.continue_trace(env, **options)
-        transaction = Sentry.start_transaction(transaction: transaction, custom_sampling_context: { env: env }, **options)
+        transaction = Sentry.continue_trace(env, established: established, **options)
+        transaction = Sentry.start_transaction(transaction: transaction, custom_sampling_context: { env: env }, established: established, **options)
         attach_queue_time(transaction, env)
         transaction
       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,8 @@ module Sentry
   class Railtie < ::Rails::Railtie
     # middlewares can't be injected after initialize
     initializer "sentry.use_rack_middleware" do |app|
+      # placed first so anything logged before CaptureExceptions shares the same trace context
+      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

@@ -118,10 +118,17 @@ def pop_scope
       end
     end
 
-    def start_transaction(transaction: nil, custom_sampling_context: {}, instrumenter: :sentry, **options)
+    def start_transaction(transaction: nil, custom_sampling_context: {}, instrumenter: :sentry, established: false, **options)
       return unless configuration.tracing_enabled?
       return unless instrumenter == configuration.instrumenter
 
+      if transaction.nil? && !options.key?(:trace_id) && established
+        # reuse the already-established trace_id instead of generating an unrelated one
+        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 = {
@@ -373,8 +380,9 @@ def get_trace_propagation_meta
       end.join("\n")
     end
 
-    def continue_trace(env, **options)
-      configure_scope { |s| s.generate_propagation_context(env) }
+    def continue_trace(env, established: false, **options)
+      # don't clobber context an earlier point in the stack already established
+      configure_scope { |s| s.generate_propagation_context(env) } unless established
 
       return nil unless configuration.tracing_enabled?
 

sentry-ruby/lib/sentry/propagation_context.rb

@@ -13,6 +13,11 @@ class PropagationContext
       "-?([01])?\\z"         # sampled
     )
 
+    # Rack env key signaling that trace context was already established earlier in
+    # the middleware stack (e.g. by +Sentry::Rails::CaptureContext+); consumed once by
+    # +Sentry::Rack::CaptureExceptions+.
+    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,16 +16,19 @@ 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
+        # consumed atomically, first, so it can't leak into later reuses of this env
+        established = env.delete(Sentry::PropagationContext::ESTABLISHED_ENV_KEY)
+
+        # make sure the current thread has a clean hub, unless it was already established
+        Sentry.clone_hub_to_current_thread unless established
 
         Sentry.with_scope do |scope|
           Sentry.with_session_tracking do
             scope.clear_breadcrumbs
             scope.set_transaction_name(env["PATH_INFO"], source: :url) if env["PATH_INFO"]
             scope.set_rack_env(env)
 
-            transaction = start_transaction(env, scope)
+            transaction = start_transaction(env, scope, established)
             scope.set_span(transaction) if transaction
 
             begin
@@ -63,16 +68,16 @@ def capture_exception(exception, env)
         end
       end
 
-      def start_transaction(env, scope)
+      def start_transaction(env, scope, established)
         options = {
           name: scope.transaction_name,
           source: scope.transaction_source,
           op: transaction_op,
           origin: SPAN_ORIGIN
         }
 
-        transaction = Sentry.continue_trace(env, **options)
-        transaction = Sentry.start_transaction(transaction: transaction, custom_sampling_context: { env: env }, **options)
+        transaction = Sentry.continue_trace(env, established: established, **options)
+        transaction = Sentry.start_transaction(transaction: transaction, custom_sampling_context: { env: env }, established: established, **options)
         attach_queue_time(transaction, env)
         transaction
       end

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|