Extract ActiveSupport::ExecutionContext out of ActiveRecord::QueryLogs

I'm working on a standardized error reporting interface
(#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 @@ def process(action, *args)
 
       @_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/railtie.rb

@@ -113,10 +113,6 @@ class Railtie < Rails::Railtie # :nodoc:
       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 @@ def perform(*)
 
     private
       def _perform_job
+        ActiveSupport::ExecutionContext[:job] = self
         run_callbacks :perform do
           perform(*arguments)
         end

activejob/lib/active_job/railtie.rb

@@ -65,10 +65,6 @@ class Railtie < Rails::Railtie # :nodoc:
       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 QueryLogs
     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 @@ def call(sql) # :nodoc:
         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 @@ def uncached_comment
           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 @@ class Railtie < Rails::Railtie # :nodoc:
           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 @@ def setup
   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 @@ def test_add_comments_to_beginning_of_query
     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 @@ def test_retrieves_comment_from_cache_when_enabled_and_set
     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 @@ def test_default_tag_behavior
   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 @@ def test_multiple_custom_tags
     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