Add support for rack.response_finished (#3681)

lib/puma/binder.rb

@@ -44,7 +44,10 @@ def initialize(log_writer, conf = Configuration.new, env: ENV)
 
         "QUERY_STRING".freeze => "",
         SERVER_SOFTWARE => PUMA_SERVER_STRING,
-        GATEWAY_INTERFACE => CGI_VER
+        GATEWAY_INTERFACE => CGI_VER,
+
+        RACK_AFTER_REPLY => nil,
+        RACK_RESPONSE_FINISHED => nil,
       }
 
       @envs = {}

lib/puma/const.rb

@@ -228,6 +228,7 @@ module Const
     RACK_INPUT = "rack.input"
     RACK_URL_SCHEME = "rack.url_scheme"
     RACK_AFTER_REPLY = "rack.after_reply"
+    RACK_RESPONSE_FINISHED = "rack.response_finished"
     PUMA_SOCKET = "puma.socket"
     PUMA_CONFIG = "puma.config"
     PUMA_PEERCERT = "puma.peercert"

lib/puma/request.rb

@@ -52,6 +52,7 @@ def handle_request(client, requests)
       io_buffer = client.io_buffer
       socket  = client.io   # io may be a MiniSSL::Socket
       app_body = nil
+      error = nil
 
       return false if closed_socket?(socket)
 
@@ -92,6 +93,7 @@ def handle_request(client, requests)
       # array, we will invoke them when the request is done.
       #
       env[RACK_AFTER_REPLY] ||= []
+      env[RACK_RESPONSE_FINISHED] ||= []
 
       begin
         if @supported_http_methods == :any || @supported_http_methods.key?(env[REQUEST_METHOD])
@@ -119,15 +121,15 @@ def handle_request(client, requests)
 
           return :async
         end
-      rescue ThreadPool::ForceShutdown => e
-        @log_writer.unknown_error e, client, "Rack app"
+      rescue ThreadPool::ForceShutdown => error
+        @log_writer.unknown_error error, client, "Rack app"
         @log_writer.log "Detected force shutdown of a thread"
 
-        status, headers, res_body = lowlevel_error(e, env, 503)
-      rescue Exception => e
-        @log_writer.unknown_error e, client, "Rack app"
+        status, headers, res_body = lowlevel_error(error, env, 503)
+      rescue Exception => error
+        @log_writer.unknown_error error, client, "Rack app"
 
-        status, headers, res_body = lowlevel_error(e, env, 500)
+        status, headers, res_body = lowlevel_error(error, env, 500)
       end
       prepare_response(status, headers, res_body, requests, client)
     ensure
@@ -144,6 +146,16 @@ def handle_request(client, requests)
           end
         end
       end
+
+      if response_finished = env[RACK_RESPONSE_FINISHED]
+        response_finished.reverse_each do |o|
+          begin
+            o.call(env, status, headers, error)
+          rescue StandardError => e
+            @log_writer.debug_error e
+          end
+        end
+      end
     end
 
     # Assembles the headers and prepares the body for actually sending the

test/test_rack_server.rb

@@ -57,7 +57,7 @@ def call(env)
 
   def setup
     @simple = lambda { |env| [200, { "x-header" => "Works" }, ["Hello"]] }
-    @server = Puma::Server.new @simple
+    @server = Puma::Server.new @simple, nil, log_writer: Puma::LogWriter.null
     @port = (@server.add_tcp_listener HOST, 0).addr[1]
     @tcp = "http://#{HOST}:#{@port}"
     @stopped = false
@@ -196,6 +196,58 @@ def test_after_reply_exception
     stop
   end
 
+  def test_rack_response_finished
+    calls = []
+
+    @server.app = lambda do |env|
+      env['rack.response_finished'] << lambda { |c_env, status, headers, error|
+        calls << 1
+        assert_same env, c_env
+        assert_equal 200, status
+        assert_instance_of Hash, headers
+        assert_nil error
+      }
+      env['rack.response_finished'] << lambda { |env, status, headers, error| calls << 2; raise "Oops" }
+      env['rack.response_finished'] << lambda { |env, status, headers, error| calls << 3 }
+      @simple.call(env)
+    end
+
+
+    @server.run
+
+    hit(["#{@tcp}/test"])
+
+    stop
+
+    assert_equal [3, 2, 1], calls
+  end
+
+  def test_rack_response_finished_on_error
+    calls = []
+
+    @server.app = lambda do |env|
+      env['rack.response_finished'] << lambda { |c_env, status, headers, error|
+        begin
+          assert_same env, c_env
+          assert_equal 500, status
+          assert_instance_of Hash, headers
+          assert_instance_of RuntimeError, error
+          assert_equal "test_rack_response_finished_on_error", error.message
+          calls << 1
+        end
+      }
+      raise "test_rack_response_finished_on_error"
+    end
+
+    @server.run
+
+    hit(["#{@tcp}/test"])
+
+    stop
+
+    assert_equal [1], calls
+  end
+
   def test_rack_body_proxy
     closed = false
     body = Rack::BodyProxy.new(["Hello"]) { closed = true }