BroadcastLogger should support TaggedLogging correctly

activesupport/lib/active_support/tagged_logging.rb

@@ -27,6 +27,30 @@ module ActiveSupport
   # it easy to stamp log lines with subdomains, request ids, and anything else
   # to aid debugging of multi-user production applications.
   module TaggedLogging
+    # This module is included in ActiveSupport::BroadcastLogger to enable
+    # broadcasting to tagged loggers with equivalent semantics.
+    module Broadcasting
+      def tagged(*tags, &block)
+        return super unless broadcasts.any? { |logger| logger.respond_to?(:tagged) }
+
+        if block_given?
+          # `tagged(...) { |logger| ... }` yields itself to the block
+          broadcasts.inject(block) do |block, logger|
+            if logger.respond_to?(:tagged)
+              proc { logger.tagged(*tags) { block.call(self) } }
+            else
+              block
+            end
+          end.call
+        else
+          # `tagged(...) returns a new logger with the tags pushed
+          self.class.new(*broadcasts.map { |logger|
+            logger.respond_to?(:tagged) ? logger.tagged(*tags) : logger
+          })
+        end
+      end
+    end
+
     module Formatter # :nodoc:
       # This method is invoked when a log event occurs.
       def call(severity, timestamp, progname, msg)
@@ -155,3 +179,5 @@ def flush
     end
   end
 end
+
+ActiveSupport::BroadcastLogger.include ActiveSupport::TaggedLogging::Broadcasting

activesupport/test/tagged_logging_test.rb

@@ -259,3 +259,68 @@ def crozz_method
     assert_equal "[tag] [1, 2, 3]\n", @output.string
   end
 end
+
+class TaggedLoggingBroadcastingTest < ActiveSupport::TestCase
+  setup do
+    @output1 = StringIO.new
+    @output2 = StringIO.new
+    @logger1 = ActiveSupport::TaggedLogging.new(Logger.new(@output1))
+    @logger2 = ActiveSupport::TaggedLogging.new(Logger.new(@output2))
+    @logger = ActiveSupport::BroadcastLogger.new(@logger1, @logger2)
+  end
+
+  test "tags with a block" do
+    @logger.tagged("BCX") { @logger.info "Funky time" }
+    assert_equal "[BCX] Funky time\n", @output1.string
+    assert_equal "[BCX] Funky time\n", @output1.string
+  end
+
+  test "yields the logger to the block" do
+    @logger.tagged("BCX") do |logger|
+      assert_equal @logger, logger
+    end
+  end
+
+  test "returns the return value of the block" do
+    return_value = @logger.tagged("BCX") do
+      "Cool story"
+    end
+    assert_equal "Cool story", return_value
+  end
+
+  test "it returns a tagged logger without a block" do
+    logger = @logger.tagged("BCX")
+    logger.info "Funky time"
+    assert_equal "[BCX] Funky time\n", @output1.string
+    assert_equal "[BCX] Funky time\n", @output2.string
+  end
+
+  test "it doesn't tag the original logger without a block" do
+    @logger.tagged("BCX")
+    @logger.info "Funky time"
+    assert_equal "Funky time\n", @output1.string
+    assert_equal "Funky time\n", @output2.string
+  end
+
+  test "it ignores non-tagged loggers with a block" do
+    plain_output = StringIO.new
+    plain_logger = ActiveSupport::Logger.new(plain_output)
+    @logger.broadcast_to(plain_logger)
+
+    @logger.tagged("BCX") { @logger.info "Block funky time" }
+    @logger.tagged("BCX").info "Chain cool story"
+    assert_equal(<<~STR, plain_output.string)
+      Block funky time
+      Chain cool story
+    STR
+  end
+
+  test "it ignores non-tagged loggers without a block" do
+    plain_output = StringIO.new
+    plain_logger = ActiveSupport::Logger.new(plain_output)
+    @logger.broadcast_to(plain_logger)
+
+    @logger.tagged("BCX").info "Cool story"
+    assert_equal "Cool story\n", plain_output.string
+  end
+end