Add support for rack.response_finished callbacks in ActionDispatch::Executor.

The executor middleware now supports deferring completion callbacks to later
in the request lifecycle by utilizing Rack's `rack.response_finished` mechanism,
when available. This enables applications to define `rack.response_finished` callbacks
that may rely on state that would be cleaned up by the executor's completion callbacks.

Co-authored-by: Hartley McGuire <[email protected]>

Author: adrianna-chang-shopify

actionpack/CHANGELOG.md

@@ -1,3 +1,13 @@
+
+*   Add support for `rack.response_finished` callbacks in ActionDispatch::Executor.
+
+    The executor middleware now supports deferring completion callbacks to later
+    in the request lifecycle by utilizing Rack's `rack.response_finished` mechanism,
+    when available. This enables applications to define `rack.response_finished` callbacks
+    that may rely on state that would be cleaned up by the executor's completion callbacks.
+
+    *Adrianna Chang*, *Hartley McGuire*
+
 *   Add JSON support to the built-in health controller.
 
     The health controller now responds to JSON requests with a structured response

actionpack/lib/action_dispatch/middleware/executor.rb

@@ -12,6 +12,10 @@ def initialize(app, executor)
 
     def call(env)
       state = @executor.run!(reset: true)
+      if response_finished = env["rack.response_finished"]
+        response_finished << -> { state.complete! }
+      end
+
       begin
         response = @app.call(env)
 
@@ -20,15 +24,21 @@ def call(env)
           @executor.error_reporter.report(error, handled: false, source: "application.action_dispatch")
         end
 
-        returned = response << ::Rack::BodyProxy.new(response.pop) { state.complete! }
+        unless response_finished
+          response << ::Rack::BodyProxy.new(response.pop) { state.complete! }
+        end
+        returned = true
+        response
       rescue Exception => error
         request = ActionDispatch::Request.new env
         backtrace_cleaner = request.get_header("action_dispatch.backtrace_cleaner")
         wrapper = ExceptionWrapper.new(backtrace_cleaner, error)
         @executor.error_reporter.report(wrapper.unwrapped_exception, handled: false, source: "application.action_dispatch")
         raise
       ensure
-        state.complete! unless returned
+        if !returned && !response_finished
+          state.complete!
+        end
       end
     end
   end

actionpack/test/dispatch/executor_test.rb

@@ -170,10 +170,44 @@ def test_handled_error_is_not_reported
     end
   end
 
+  def test_complete_callbacks_are_called_on_rack_response_finished
+    completed = false
+    executor.to_complete { completed = true }
+
+    env = { "rack.response_finished" => [] }
+    call_and_return_body(env)
+
+    assert_not completed
+
+    assert_equal 1, env["rack.response_finished"].size
+    env["rack.response_finished"].first.call
+
+    assert completed
+  end
+
+  def test_complete_callbacks_are_called_once_on_rack_response_finished_when_exception_is_raised
+    completed_count = 0
+    executor.to_complete { completed_count += 1 }
+
+    env = { "rack.response_finished" => [] }
+
+    begin
+      call_and_return_body(env) do
+        raise "error"
+      end
+    rescue
+    end
+
+    assert_equal 1, env["rack.response_finished"].size
+    env["rack.response_finished"].first.call
+
+    assert_equal 1, completed_count
+  end
+
   private
-    def call_and_return_body(&block)
+    def call_and_return_body(env = {}, &block)
       app = block || proc { [200, {}, []] }
-      env = Rack::MockRequest.env_for("", {})
+      env = Rack::MockRequest.env_for("", env)
       _, _, body = middleware(app).call(env)
       body
     end