Extract ActiveSupport::ExecutionContext out of ActiveRecord::QueryLogs

I'm working on a standardized error reporting interface
(https://github.com/rails/rails/issues/43472) and it has the same need
for a `context` than Active Record's query logs.

Just like query logs, when reporting an error you want to know what the
current controller or job is, etc.

So by extracting it we can allow both API to use the same store.

actionpack/lib/abstract_controller/base.rb

@@ -148,6 +148,7 @@ module AbstractController
 
       @_response_body = nil
 
+      ActiveSupport::ExecutionContext[:controller] = self
       process_action(action_name, *args)
     end
 

actionpack/lib/action_controller.rb

@@ -37,7 +37,6 @@ module ActionController
     autoload :Logging
     autoload :MimeResponds
     autoload :ParamsWrapper
-    autoload :QueryTags
     autoload :Redirecting
     autoload :Renderers
     autoload :Rendering

actionpack/lib/action_controller/metal/query_tags.rb

@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-
-module ActionController
-  module QueryTags # :nodoc:
-    extend ActiveSupport::Concern
-
-    included do
-      around_action :expose_controller_to_query_logs
-    end
-
-    private
-      def expose_controller_to_query_logs(&block)
-        ActiveRecord::QueryLogs.set_context(controller: self, &block)
-      end
-  end
-end

actionpack/lib/action_controller/railtie.rb

@@ -113,10 +113,6 @@ module ActionController
       if query_logs_tags_enabled
         app.config.active_record.query_log_tags += [:controller, :action]
 
-        ActiveSupport.on_load(:action_controller) do
-          include ActionController::QueryTags
-        end
-
         ActiveSupport.on_load(:active_record) do
           ActiveRecord::QueryLogs.taggings.merge!(
             controller:            ->(context) { context[:controller]&.controller_name },

activejob/lib/active_job/execution.rb

@@ -54,6 +54,7 @@ module ActiveJob
 
     private
       def _perform_job
+        ActiveSupport::ExecutionContext[:job] = self
         run_callbacks :perform do
           perform(*arguments)
         end

activejob/lib/active_job/query_tags.rb

@@ -1,16 +0,0 @@
-# frozen_string_literal: true
-
-module ActiveJob
-  module QueryTags # :nodoc:
-    extend ActiveSupport::Concern
-
-    included do
-      around_perform :expose_job_to_query_logs
-    end
-
-    private
-      def expose_job_to_query_logs(&block)
-        ActiveRecord::QueryLogs.set_context(job: self, &block)
-      end
-  end
-end

activejob/lib/active_job/railtie.rb

@@ -65,10 +65,6 @@ module ActiveJob
       if query_logs_tags_enabled
         app.config.active_record.query_log_tags << :job
 
-        ActiveSupport.on_load(:active_job) do
-          include ActiveJob::QueryTags
-        end
-
         ActiveSupport.on_load(:active_record) do
           ActiveRecord::QueryLogs.taggings[:job] = ->(context) { context[:job].class.name if context[:job] }
         end

activerecord/lib/active_record/query_logs.rb

@@ -44,17 +44,17 @@ module ActiveRecord
   #    ]
   #    ActiveRecord::QueryLogs.tags = tags
   #
-  # The QueryLogs +context+ can be manipulated via the +set_context+ method.
+  # The QueryLogs +context+ can be manipulated via the +ActiveSupport::ExecutionContext.set+ method.
   #
   # Temporary updates limited to the execution of a block:
   #
-  #    ActiveRecord::QueryLogs.set_context(foo: Bar.new) do
+  #    ActiveSupport::ExecutionContext.set(foo: Bar.new) do
   #      posts = Post.all
   #    end
   #
   # Direct updates to a context value:
   #
-  #    ActiveRecord::QueryLogs.set_context(foo: Bar.new)
+  #    ActiveSupport::ExecutionContext[:foo] = Bar.new
   #
   # Tag comments can be prepended to the query:
   #
@@ -76,30 +76,6 @@ module ActiveRecord
     thread_mattr_accessor :cached_comment, instance_accessor: false
 
     class << self
-      # Updates the context used to construct tags in the SQL comment.
-      # If a block is given, it resets the provided keys to their
-      # previous value once the block exits.
-      def set_context(**options)
-        options.symbolize_keys!
-
-        keys = options.keys
-        previous_context = keys.zip(context.values_at(*keys)).to_h
-        context.merge!(options)
-        self.cached_comment = nil
-        if block_given?
-          begin
-            yield
-          ensure
-            context.merge!(previous_context)
-            self.cached_comment = nil
-          end
-        end
-      end
-
-      def clear_context # :nodoc:
-        context.clear
-      end
-
       def call(sql) # :nodoc:
         if prepend_comment
           "#{self.comment} #{sql}"
@@ -108,6 +84,12 @@ module ActiveRecord
         end.strip
       end
 
+      def clear_cache # :nodoc:
+        self.cached_comment = nil
+      end
+
+      ActiveSupport::ExecutionContext.after_change { ActiveRecord::QueryLogs.clear_cache }
+
       private
         # Returns an SQL comment +String+ containing the query log tags.
         # Sets and returns a cached comment if <tt>cache_query_log_tags</tt> is +true+.
@@ -126,15 +108,13 @@ module ActiveRecord
           end
         end
 
-        def context
-          Thread.current[:active_record_query_log_tags_context] ||= {}
-        end
-
         def escape_sql_comment(content)
           content.to_s.gsub(%r{ (/ (?: | \g<1>) \*) \+? \s* | \s* (\* (?: | \g<2>) /) }x, "")
         end
 
         def tag_content
+          context = ActiveSupport::ExecutionContext.to_h
+
           tags.flat_map { |i| [*i] }.filter_map do |tag|
             key, handler = tag
             handler ||= taggings[key]

activerecord/lib/active_record/railtie.rb

@@ -376,10 +376,6 @@ To keep using the current cache store, you can turn off cache versioning entirel
           if app.config.active_record.cache_query_log_tags
             ActiveRecord::QueryLogs.cache_query_log_tags = true
           end
-
-          app.reloader.before_class_unload { ActiveRecord::QueryLogs.clear_context }
-          app.executor.to_run              { ActiveRecord::QueryLogs.clear_context }
-          app.executor.to_complete         { ActiveRecord::QueryLogs.clear_context }
         end
       end
     end

activerecord/test/cases/query_logs_test.rb

@@ -11,13 +11,14 @@ class QueryLogsTest < ActiveRecord::TestCase
   }
 
   def setup
-    # QueryLogs context is automatically reset in Rails app via an executor hooks set in railtie
+    # ActiveSupport::ExecutionContext context is automatically reset in Rails app via an executor hooks set in railtie
     # But not in Active Record's own test suite.
-    ActiveRecord::QueryLogs.clear_context
+    ActiveSupport::ExecutionContext.clear
 
     # Enable the query tags logging
     @original_transformers = ActiveRecord.query_transformers
     @original_prepend = ActiveRecord::QueryLogs.prepend_comment
+    @original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord.query_transformers += [ActiveRecord::QueryLogs]
     ActiveRecord::QueryLogs.prepend_comment = false
     ActiveRecord::QueryLogs.cache_query_log_tags = false
@@ -27,10 +28,13 @@ class QueryLogsTest < ActiveRecord::TestCase
   def teardown
     ActiveRecord.query_transformers = @original_transformers
     ActiveRecord::QueryLogs.prepend_comment = @original_prepend
-    ActiveRecord::QueryLogs.tags = []
-    # QueryLogs context is automatically reset in Rails app via an executor hooks set in railtie
+    ActiveRecord::QueryLogs.tags = @original_tags
+    ActiveRecord::QueryLogs.prepend_comment = false
+    ActiveRecord::QueryLogs.cache_query_log_tags = false
+    ActiveRecord::QueryLogs.cached_comment = nil
+    # ActiveSupport::ExecutionContext context is automatically reset in Rails app via an executor hooks set in railtie
     # But not in Active Record's own test suite.
-    ActiveRecord::QueryLogs.clear_context
+    ActiveSupport::ExecutionContext.clear
   end
 
   def test_escaping_good_comment
@@ -57,8 +61,6 @@ class QueryLogsTest < ActiveRecord::TestCase
     assert_sql(%r{/\*application:active_record\*/ select id from posts$}) do
       ActiveRecord::Base.connection.execute "select id from posts"
     end
-  ensure
-    ActiveRecord::QueryLogs.prepend_comment = nil
   end
 
   def test_exists_is_commented
@@ -112,41 +114,24 @@ class QueryLogsTest < ActiveRecord::TestCase
     ActiveRecord::QueryLogs.stub(:cached_comment, "/*cached_comment*/") do
       assert_equal "/*cached_comment*/", ActiveRecord::QueryLogs.call("")
     end
-  ensure
-    ActiveRecord::QueryLogs.cached_comment = nil
-    ActiveRecord::QueryLogs.cache_query_log_tags = false
   end
 
   def test_resets_cache_on_context_update
     ActiveRecord::QueryLogs.cache_query_log_tags = true
-    ActiveRecord::QueryLogs.set_context(temporary: "value")
+    ActiveSupport::ExecutionContext[:temporary] = "value"
     ActiveRecord::QueryLogs.tags = [ temporary_tag: ->(context) { context[:temporary] } ]
 
     assert_equal "/*temporary_tag:value*/", ActiveRecord::QueryLogs.call("")
 
-    ActiveRecord::QueryLogs.set_context(temporary: "new_value")
+    ActiveSupport::ExecutionContext[:temporary] = "new_value"
 
     assert_nil ActiveRecord::QueryLogs.cached_comment
     assert_equal "/*temporary_tag:new_value*/", ActiveRecord::QueryLogs.call("")
-  ensure
-    ActiveRecord::QueryLogs.cached_comment = nil
-    ActiveRecord::QueryLogs.cache_query_log_tags = false
-  end
-
-  def test_ensure_context_has_symbol_keys
-    ActiveRecord::QueryLogs.tags = [ new_key: ->(context) { context[:symbol_key] } ]
-    ActiveRecord::QueryLogs.set_context("symbol_key" => "symbolized")
-
-    assert_sql(%r{/\*new_key:symbolized}) do
-      Dashboard.first
-    end
-  ensure
-    ActiveRecord::QueryLogs.set_context(application_name: nil)
   end
 
   def test_default_tag_behavior
     ActiveRecord::QueryLogs.tags = [:application, :foo]
-    ActiveRecord::QueryLogs.set_context(foo: "bar") do
+    ActiveSupport::ExecutionContext.set(foo: "bar") do
       assert_sql(%r{/\*application:active_record,foo:bar\*/}) do
         Dashboard.first
       end
@@ -157,39 +142,29 @@ class QueryLogsTest < ActiveRecord::TestCase
   end
 
   def test_empty_comments_are_not_added
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ empty: -> { nil } ]
     assert_sql(%r{select id from posts$}) do
       ActiveRecord::Base.connection.execute "select id from posts"
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_custom_basic_tags
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ :application, { custom_string: "test content" } ]
 
     assert_sql(%r{/\*application:active_record,custom_string:test content\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_custom_proc_tags
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [ :application, { custom_proc: -> { "test content" } } ]
 
     assert_sql(%r{/\*application:active_record,custom_proc:test content\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_multiple_custom_tags
-    original_tags = ActiveRecord::QueryLogs.tags
     ActiveRecord::QueryLogs.tags = [
       :application,
       { custom_proc: -> { "test content" }, another_proc: -> { "more test content" } },
@@ -198,34 +173,14 @@ class QueryLogsTest < ActiveRecord::TestCase
     assert_sql(%r{/\*application:active_record,custom_proc:test content,another_proc:more test content\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 
   def test_custom_proc_context_tags
-    original_tags = ActiveRecord::QueryLogs.tags
-    ActiveRecord::QueryLogs.set_context(foo: "bar")
+    ActiveSupport::ExecutionContext[:foo] = "bar"
     ActiveRecord::QueryLogs.tags = [ :application, { custom_context_proc: ->(context) { context[:foo] } } ]
 
     assert_sql(%r{/\*application:active_record,custom_context_proc:bar\*/$}) do
       Dashboard.first
     end
-  ensure
-    ActiveRecord::QueryLogs.set_context(foo: nil)
-    ActiveRecord::QueryLogs.tags = original_tags
-  end
-
-  def test_set_context_restore_state
-    original_tags = ActiveRecord::QueryLogs.tags
-    ActiveRecord::QueryLogs.tags = [foo: ->(context) { context[:foo] }]
-    ActiveRecord::QueryLogs.set_context(foo: "bar") do
-      assert_sql(%r{/\*foo:bar\*/$}) { Dashboard.first }
-      ActiveRecord::QueryLogs.set_context(foo: "plop") do
-        assert_sql(%r{/\*foo:plop\*/$}) { Dashboard.first }
-      end
-      assert_sql(%r{/\*foo:bar\*/$}) { Dashboard.first }
-    end
-  ensure
-    ActiveRecord::QueryLogs.tags = original_tags
   end
 end

activesupport/lib/active_support.rb

@@ -40,6 +40,7 @@ module ActiveSupport
   autoload :CurrentAttributes
   autoload :Dependencies
   autoload :DescendantsTracker
+  autoload :ExecutionContext
   autoload :ExecutionWrapper
   autoload :Executor
   autoload :FileUpdateChecker

activesupport/lib/active_support/execution_context.rb

@@ -0,0 +1,53 @@
+# frozen_string_literal: true
+
+module ActiveSupport
+  module ExecutionContext # :nodoc:
+    @after_change_callbacks = []
+    class << self
+      def after_change(&block)
+        @after_change_callbacks << block
+      end
+
+      # Updates the execution context. If a block is given, it resets the provided keys to their
+      # previous value once the block exits.
+      def set(**options)
+        options.symbolize_keys!
+        keys = options.keys
+
+        store = self.store
+
+        previous_context = keys.zip(store.values_at(*keys)).to_h
+
+        store.merge!(options)
+        @after_change_callbacks.each(&:call)
+
+        if block_given?
+          begin
+            yield
+          ensure
+            store.merge!(previous_context)
+            @after_change_callbacks.each(&:call)
+          end
+        end
+      end
+
+      def []=(key, value)
+        store[key.to_sym] = value
+        @after_change_callbacks.each(&:call)
+      end
+
+      def to_h
+        store.dup
+      end
+
+      def clear
+        store.clear
+      end
+
+      private
+        def store
+          Thread.current[:active_support_execution_context] ||= {}
+        end
+    end
+  end
+end

activesupport/lib/active_support/execution_context/test_helper.rb

@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+module ActiveSupport::ExecutionContext::TestHelper # :nodoc:
+  def before_setup
+    ActiveSupport::ExecutionContext.clear
+    super
+  end
+
+  def after_teardown
+    super
+    ActiveSupport::ExecutionContext.clear
+  end
+end

activesupport/lib/active_support/railtie.rb

@@ -27,6 +27,12 @@ module ActiveSupport
       end
     end
 
+    initializer "active_support.reset_execution_context" do |app|
+      app.reloader.before_class_unload { ActiveSupport::ExecutionContext.clear }
+      app.executor.to_run              { ActiveSupport::ExecutionContext.clear }
+      app.executor.to_complete         { ActiveSupport::ExecutionContext.clear }
+    end
+
     initializer "active_support.reset_all_current_attributes_instances" do |app|
       executor_around_test_case = app.config.active_support.delete(:executor_around_test_case)
 
@@ -41,6 +47,9 @@ module ActiveSupport
         else
           require "active_support/current_attributes/test_helper"
           include ActiveSupport::CurrentAttributes::TestHelper
+
+          require "active_support/execution_context/test_helper"
+          include ActiveSupport::ExecutionContext::TestHelper
         end
       end
     end

activesupport/test/current_attributes_test.rb

@@ -1,17 +1,12 @@
 # frozen_string_literal: true
 
 require_relative "abstract_unit"
+require "active_support/current_attributes/test_helper"
 
 class CurrentAttributesTest < ActiveSupport::TestCase
   # CurrentAttributes is automatically reset in Rails app via executor hooks set in railtie
   # But not in Active Support's own test suite.
-  setup do
-    ActiveSupport::CurrentAttributes.reset_all
-  end
-
-  teardown do
-    ActiveSupport::CurrentAttributes.reset_all
-  end
+  include ActiveSupport::CurrentAttributes::TestHelper
 
   Person = Struct.new(:id, :name, :time_zone)
 

activesupport/test/execution_context_test.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require_relative "abstract_unit"
+require "active_support/execution_context/test_helper"
+
+class ExecutionContextTest < ActiveSupport::TestCase
+  # ExecutionContext is automatically reset in Rails app via executor hooks set in railtie
+  # But not in Active Support's own test suite.
+  include ActiveSupport::ExecutionContext::TestHelper
+
+  test "#set restore the modified keys when the block exits" do
+    assert_nil ActiveSupport::ExecutionContext.to_h[:foo]
+    ActiveSupport::ExecutionContext.set(foo: "bar") do
+      assert_equal "bar", ActiveSupport::ExecutionContext.to_h[:foo]
+      ActiveSupport::ExecutionContext.set(foo: "plop") do
+        assert_equal "plop", ActiveSupport::ExecutionContext.to_h[:foo]
+      end
+      assert_equal "bar", ActiveSupport::ExecutionContext.to_h[:foo]
+
+      ActiveSupport::ExecutionContext[:direct_assignment] = "present"
+      ActiveSupport::ExecutionContext.set(multi_assignment: "present")
+    end
+
+    assert_nil ActiveSupport::ExecutionContext.to_h[:foo]
+
+    assert_equal "present", ActiveSupport::ExecutionContext.to_h[:direct_assignment]
+    assert_equal "present", ActiveSupport::ExecutionContext.to_h[:multi_assignment]
+  end
+
+  test "#set coerce keys to symbol" do
+    ActiveSupport::ExecutionContext.set("foo" => "bar") do
+      assert_equal "bar", ActiveSupport::ExecutionContext.to_h[:foo]
+    end
+  end
+
+  test "#[]= coerce keys to symbol" do
+    ActiveSupport::ExecutionContext["symbol_key"] = "symbolized"
+    assert_equal "symbolized", ActiveSupport::ExecutionContext.to_h[:symbol_key]
+  end
+
+  test "#to_h returns a copy of the context" do
+    ActiveSupport::ExecutionContext[:foo] = 42
+    context = ActiveSupport::ExecutionContext.to_h
+    context[:foo] = 43
+    assert_equal 42, ActiveSupport::ExecutionContext.to_h[:foo]
+  end
+end