Directly support stateful executor hooks

Also, make sure to call the +complete+ hooks if +run+ fails.

activerecord/lib/active_record/query_cache.rb

@@ -23,23 +23,27 @@ def uncached(&block)
       end
     end
 
-    def self.install_executor_hooks(executor = ActiveSupport::Executor)
-      executor.to_run do
-        connection    = ActiveRecord::Base.connection
-        enabled       = connection.query_cache_enabled
-        connection_id = ActiveRecord::Base.connection_id
-        connection.enable_query_cache!
+    def self.run
+      connection    = ActiveRecord::Base.connection
+      enabled       = connection.query_cache_enabled
+      connection_id = ActiveRecord::Base.connection_id
+      connection.enable_query_cache!
 
-        @restore_query_cache_settings = lambda do
-          ActiveRecord::Base.connection_id = connection_id
-          ActiveRecord::Base.connection.clear_query_cache
-          ActiveRecord::Base.connection.disable_query_cache! unless enabled
-        end
-      end
+      [enabled, connection_id]
+    end
 
-      executor.to_complete do
-        @restore_query_cache_settings.call if defined?(@restore_query_cache_settings)
+    def self.complete(state)
+      enabled, connection_id = state
 
+      ActiveRecord::Base.connection_id = connection_id
+      ActiveRecord::Base.connection.clear_query_cache
+      ActiveRecord::Base.connection.disable_query_cache! unless enabled
+    end
+
+    def self.install_executor_hooks(executor = ActiveSupport::Executor)
+      executor.register_hook(self)
+
+      executor.to_complete do
         # FIXME: This should be skipped when env['rack.test']
         ActiveRecord::Base.clear_active_connections!
       end

activesupport/lib/active_support/execution_wrapper.rb

@@ -19,6 +19,32 @@ def self.to_complete(*args, &block)
       set_callback(:complete, *args, &block)
     end
 
+    # Register an object to be invoked during both the +run+ and
+    # +complete+ steps.
+    #
+    # +hook.complete+ will be passed the value returned from +hook.run+,
+    # and will only be invoked if +run+ has previously been called.
+    # (Mostly, this means it won't be invoked if an exception occurs in
+    # a preceding +to_run+ block; all ordinary +to_complete+ blocks are
+    # invoked in that situation.)
+    def self.register_hook(hook, outer: false)
+      if outer
+        run_args = [prepend: true]
+        complete_args = [:after]
+      else
+        run_args = complete_args = []
+      end
+
+      to_run(*run_args) do
+        hook_state[hook] = hook.run
+      end
+      to_complete(*complete_args) do
+        if hook_state.key?(hook)
+          hook.complete hook_state[hook]
+        end
+      end
+    end
+
     # Run this execution.
     #
     # Returns an instance, whose +complete!+ method *must* be invoked
@@ -29,19 +55,27 @@ def self.run!
       if active?
         Null
       else
-        new.tap(&:run!)
+        new.tap do |instance|
+          success = nil
+          begin
+            instance.run!
+            success = true
+          ensure
+            instance.complete! unless success
+          end
+        end
       end
     end
 
     # Perform the work in the supplied block as an execution.
     def self.wrap
       return yield if active?
 
-      state = run!
+      instance = run!
       begin
         yield
       ensure
-        state.complete!
+        instance.complete!
       end
     end
 
@@ -74,5 +108,10 @@ def complete!
     ensure
       self.class.active.delete Thread.current
     end
+
+    private
+      def hook_state
+        @_hook_state ||= {}
+      end
   end
 end

activesupport/lib/active_support/reloader.rb

@@ -43,7 +43,13 @@ def self.after_class_unload(*args, &block)
     # Initiate a manual reload
     def self.reload!
       executor.wrap do
-        new.tap(&:run!).complete!
+        new.tap do |instance|
+          begin
+            instance.run!
+          ensure
+            instance.complete!
+          end
+        end
       end
       prepare!
     end

activesupport/test/executor_test.rb

@@ -1,6 +1,9 @@
 require 'abstract_unit'
 
 class ExecutorTest < ActiveSupport::TestCase
+  class DummyError < RuntimeError
+  end
+
   def test_wrap_invokes_callbacks
     called = []
     executor.to_run { called << :run }
@@ -35,6 +38,20 @@ def test_separated_calls_invoke_callbacks
     assert_equal [:run, :body, :complete], called
   end
 
+  def test_exceptions_unwind
+    called = []
+    executor.to_run { called << :run_1 }
+    executor.to_run { raise DummyError }
+    executor.to_run { called << :run_2 }
+    executor.to_complete { called << :complete }
+
+    assert_raises(DummyError) do
+      executor.wrap { called << :body }
+    end
+
+    assert_equal [:run_1, :complete], called
+  end
+
   def test_avoids_double_wrapping
     called = []
     executor.to_run { called << :run }
@@ -51,6 +68,96 @@ def test_avoids_double_wrapping
     assert_equal [:run, :early, :body, :late, :complete], called
   end
 
+  def test_hooks_carry_state
+    supplied_state = :none
+
+    hook = Class.new do
+      define_method(:run) do
+        :some_state
+      end
+
+      define_method(:complete) do |state|
+        supplied_state = state
+      end
+    end.new
+
+    executor.register_hook(hook)
+
+    executor.wrap { }
+
+    assert_equal :some_state, supplied_state
+  end
+
+  def test_nil_state_is_sufficient
+    supplied_state = :none
+
+    hook = Class.new do
+      define_method(:run) do
+        nil
+      end
+
+      define_method(:complete) do |state|
+        supplied_state = state
+      end
+    end.new
+
+    executor.register_hook(hook)
+
+    executor.wrap { }
+
+    assert_equal nil, supplied_state
+  end
+
+  def test_exception_skips_uninvoked_hook
+    supplied_state = :none
+
+    hook = Class.new do
+      define_method(:run) do
+        :some_state
+      end
+
+      define_method(:complete) do |state|
+        supplied_state = state
+      end
+    end.new
+
+    executor.to_run do
+      raise DummyError
+    end
+    executor.register_hook(hook)
+
+    assert_raises(DummyError) do
+      executor.wrap { }
+    end
+
+    assert_equal :none, supplied_state
+  end
+
+  def test_exception_unwinds_invoked_hook
+    supplied_state = :none
+
+    hook = Class.new do
+      define_method(:run) do
+        :some_state
+      end
+
+      define_method(:complete) do |state|
+        supplied_state = state
+      end
+    end.new
+
+    executor.register_hook(hook)
+    executor.to_run do
+      raise DummyError
+    end
+
+    assert_raises(DummyError) do
+      executor.wrap { }
+    end
+
+    assert_equal :some_state, supplied_state
+  end
+
   def test_separate_classes_can_wrap
     other_executor = Class.new(ActiveSupport::Executor)
 

railties/lib/rails/application/finisher.rb

@@ -62,18 +62,36 @@ module Finisher
         ActiveSupport.run_load_hooks(:after_initialize, self)
       end
 
+      class MutexHook
+        def initialize(mutex = Mutex.new)
+          @mutex = mutex
+        end
+
+        def run
+          @mutex.lock
+        end
+
+        def complete(_state)
+          @mutex.unlock
+        end
+      end
+
+      module InterlockHook
+        def self.run
+          ActiveSupport::Dependencies.interlock.start_running
+        end
+
+        def self.complete(_state)
+          ActiveSupport::Dependencies.interlock.done_running
+        end
+      end
+
       initializer :configure_executor_for_concurrency do |app|
         if config.allow_concurrency == false
           # User has explicitly opted out of concurrent request
           # handling: presumably their code is not threadsafe
 
-          mutex = Mutex.new
-          app.executor.to_run(prepend: true) do
-            mutex.lock
-          end
-          app.executor.to_complete(:after) do
-            mutex.unlock
-          end
+          app.executor.register_hook(MutexHook.new, outer: true)
 
         elsif config.allow_concurrency == :unsafe
           # Do nothing, even if we know this is dangerous. This is the
@@ -86,12 +104,7 @@ module Finisher
             # Without cache_classes + eager_load, the load interlock
             # is required for proper operation
 
-            app.executor.to_run(prepend: true) do
-              ActiveSupport::Dependencies.interlock.start_running
-            end
-            app.executor.to_complete(:after) do
-              ActiveSupport::Dependencies.interlock.done_running
-            end
+            app.executor.register_hook(InterlockHook, outer: true)
           end
         end
       end