From 5faeb7007d2d85b2ea01662f65b382404cca1a97 Mon Sep 17 00:00:00 2001
From: KJ Tsanaktsidis <ktsanaktsidis@zendesk.com>
Date: Wed, 7 Feb 2024 14:05:08 +1100
Subject: [PATCH] Don't pop logger tags in Rails::Rack::Logger until request is
 finished

At the moment, Rails::Rack::Logger tags the logger (if it's
ActiveSupport::TaggedLogging) for the duration of the @app.call, but
only fires the request.action_dispatch event later, on body close. That
means anything logged in request.action_dispatch handlers won't have the
same tags as the rest of the request.

Fix this by deferring the popping of tags into
finish_request_instrumentation, in the same way that finishing the
instrumentation handle is deferred.
---
 railties/CHANGELOG.md             |  4 ++++
 railties/lib/rails/rack/logger.rb | 16 +++++++++-------
 railties/test/rack_logger_test.rb | 26 ++++++++++++++++++++++++++
 3 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/railties/CHANGELOG.md b/railties/CHANGELOG.md
index b616ab7501..7df71d6480 100644
--- a/railties/CHANGELOG.md
+++ b/railties/CHANGELOG.md
@@ -1,3 +1,7 @@
+*   Ensure logger tags configured with `config.log_tags` are still active in `request.action_dispatch` handlers
+
+    *KJ Tsanaktsidis*
+
 *   Setup jemalloc in the default Dockerfile for memory optimization.
 
     *Matt Almeida*, *Jean Boussier*
diff --git a/railties/lib/rails/rack/logger.rb b/railties/lib/rails/rack/logger.rb
index 26324e4bb8..1b561c4f1a 100644
--- a/railties/lib/rails/rack/logger.rb
+++ b/railties/lib/rails/rack/logger.rb
@@ -20,22 +20,23 @@ def initialize(app, taggers = nil)
       def call(env)
         request = ActionDispatch::Request.new(env)
 
-        if logger.respond_to?(:tagged)
-          logger.tagged(*compute_tags(request)) { call_app(request, env) }
+        logger_tag_pop_count = if logger.respond_to?(:push_tags)
+          logger.push_tags(*compute_tags(request)).size
         else
-          call_app(request, env)
+          0
         end
+        call_app(request, env, logger_tag_pop_count)
       end
 
       private
-        def call_app(request, env) # :doc:
+        def call_app(request, env, logger_tag_pop_count) # :doc:
           instrumenter = ActiveSupport::Notifications.instrumenter
           handle = instrumenter.build_handle("request.action_dispatch", { request: request })
           handle.start
 
           logger.info { started_request_message(request) }
           status, headers, body = response = @app.call(env)
-          body = ::Rack::BodyProxy.new(body) { finish_request_instrumentation(handle) }
+          body = ::Rack::BodyProxy.new(body) { finish_request_instrumentation(handle, logger_tag_pop_count) }
 
           if response.frozen?
             [status, headers, body]
@@ -44,7 +45,7 @@ def call_app(request, env) # :doc:
             response
           end
         rescue Exception
-          finish_request_instrumentation(handle)
+          finish_request_instrumentation(handle, logger_tag_pop_count)
           raise
         end
 
@@ -74,8 +75,9 @@ def logger
           Rails.logger
         end
 
-        def finish_request_instrumentation(handle)
+        def finish_request_instrumentation(handle, logger_tag_pop_count)
           handle.finish
+          logger.pop_tags(logger_tag_pop_count) if logger.respond_to?(:pop_tags) && logger_tag_pop_count > 0
           ActiveSupport::LogSubscriber.flush_all!
         end
     end
diff --git a/railties/test/rack_logger_test.rb b/railties/test/rack_logger_test.rb
index 8f9b699177..290a9c91f0 100644
--- a/railties/test/rack_logger_test.rb
+++ b/railties/test/rack_logger_test.rb
@@ -119,6 +119,32 @@ def test_logger_is_flushed_after_request_finished
       ensure
         @notifier.unsubscribe block_sub
       end
+
+      def test_logger_pushes_tags
+        @logger = ActiveSupport::TaggedLogging.new(@logger)
+        set_logger(@logger)
+
+        taggers = ["tag1", ->(_req) { "tag2" }]
+        logger_middleware = TestLogger.new(@logger, taggers: taggers) do
+          # We can't really assert on logging something with the tags, because the MockLogger implementation
+          # does not call the formatter (which is responsible for appending the tags)
+          assert_equal(["tag1", "tag2"], @logger.formatter.current_tags)
+        end
+        block_sub = @notifier.subscribe "request.action_dispatch" do |_event|
+          assert_equal(["tag1", "tag2"], @logger.formatter.current_tags)
+        end
+
+        # Call the app - it should log the inside app message
+        response_body = logger_middleware.call("REQUEST_METHOD" => "GET").last
+        # The tags should still be open as long as the request body isn't closed
+        assert_equal(["tag1", "tag2"], @logger.formatter.current_tags)
+        # And now should fire the request.action_dispatch event and call the event handler
+        response_body.close
+        # And it should also clear the tag stack.
+        assert_equal([], @logger.formatter.current_tags)
+      ensure
+        @notifier.unsubscribe block_sub
+      end
     end
   end
 end