Add logger.push_tags and .pop_tags to complement logger.tagged

Avoid memory leak from unflushed logs on other threads leaving tags behind.

Conflicts:
	activesupport/CHANGELOG.md
	activesupport/lib/active_support/tagged_logging.rb
	activesupport/test/tagged_logging_test.rb

3 files changed, 69 insertions, 18 deletions

diff --git a/activesupport/CHANGELOG.md b/activesupport/CHANGELOG.md
index 04f8568e4d..76a232665f 100644
--- a/activesupport/CHANGELOG.md
+++ b/activesupport/CHANGELOG.md
@@ -1,5 +1,19 @@
 ## Rails 3.2.9 (unreleased)
 
+*   Add logger.push_tags and .pop_tags to complement logger.tagged:
+
+        class Job
+          def before
+            Rails.logger.push_tags :jobs, self.class.name
+          end
+
+          def after
+            Rails.logger.pop_tags 2
+          end
+        end
+
+    *Jeremy Kemper*
+
 *  Add %:z and %::z format string support to ActiveSupport::TimeWithZone#strftime. [fixes #6962] *kennyj*
 
 ## Rails 3.2.8 (Aug 9, 2012) ##
diff --git a/activesupport/lib/active_support/tagged_logging.rb b/activesupport/lib/active_support/tagged_logging.rb
index 6af87e85e6..cb5c775a22 100644
--- a/activesupport/lib/active_support/tagged_logging.rb
+++ b/activesupport/lib/active_support/tagged_logging.rb
@@ -15,16 +15,27 @@ module ActiveSupport
   class TaggedLogging
     def initialize(logger)
       @logger = logger
-      @tags   = Hash.new { |h,k| h[k] = [] }
     end
 
-    def tagged(*new_tags)
-      tags     = current_tags
-      new_tags = Array.wrap(new_tags).flatten.reject(&:blank?)
-      tags.concat new_tags
-      yield
+    def tagged(*tags)
+      new_tags = push_tags *tags
+      yield self
     ensure
-      new_tags.size.times { tags.pop }
+      pop_tags new_tags.size
+    end
+
+    def push_tags(*tags)
+      tags.flatten.reject(&:blank?).tap do |new_tags|
+        current_tags.concat new_tags
+      end
+    end
+
+    def pop_tags(size = 1)
+      current_tags.pop size
+    end
+
+    def clear_tags!
+      current_tags.clear
     end
 
     def silence(temporary_level = Logger::ERROR, &block)
@@ -46,7 +57,7 @@ module ActiveSupport
     end
 
     def flush
-      @tags.delete(Thread.current)
+      clear_tags!
       @logger.flush if @logger.respond_to?(:flush)
     end
 
@@ -54,17 +65,16 @@ module ActiveSupport
       @logger.send(method, *args)
     end
 
-    protected
-
-    def tags_text
-      tags = current_tags
-      if tags.any?
-        tags.collect { |tag| "[#{tag}]" }.join(" ") + " "
+    private
+      def tags_text
+        tags = current_tags
+        if tags.any?
+          tags.collect { |tag| "[#{tag}] " }.join
+        end
       end
-    end
 
-    def current_tags
-      @tags[Thread.current]
+      def current_tags
+        Thread.current[:activesupport_tagged_logging_tags] ||= []
+      end
     end
-  end
 end
diff --git a/activesupport/test/tagged_logging_test.rb b/activesupport/test/tagged_logging_test.rb
index c838c073e6..7253eb1222 100644
--- a/activesupport/test/tagged_logging_test.rb
+++ b/activesupport/test/tagged_logging_test.rb
@@ -29,6 +29,33 @@ class TaggedLoggingTest < ActiveSupport::TestCase
     assert_equal "[BCX] [Jason] [New] Funky time\n", @output.string
   end
 
+  test "tagged are flattened" do
+    @logger.tagged("BCX", %w(Jason New)) { @logger.info "Funky time" }
+    assert_equal "[BCX] [Jason] [New] Funky time\n", @output.string
+  end
+
+  test "push and pop tags directly" do
+    assert_equal %w(A B C), @logger.push_tags('A', ['B', '  ', ['C']])
+    @logger.info 'a'
+    assert_equal %w(C), @logger.pop_tags
+    @logger.info 'b'
+    assert_equal %w(B), @logger.pop_tags(1)
+    @logger.info 'c'
+    assert_equal [], @logger.clear_tags!
+    @logger.info 'd'
+    assert_equal "[A] [B] [C] a\n[A] [B] b\n[A] c\nd\n", @output.string
+  end
+
+  test "does not strip message content" do
+    @logger.info "  Hello"
+    assert_equal "  Hello\n", @output.string
+  end
+
+  test "provides access to the logger instance" do
+    @logger.tagged("BCX") { |logger| logger.info "Funky time" }
+    assert_equal "[BCX] Funky time\n", @output.string
+  end
+
   test "tagged once with blank and nil" do
     @logger.tagged(nil, "", "New") { @logger.info "Funky time" }
     assert_equal "[New] Funky time\n", @output.string