Implement rack.response_finished

Fix: #97
Ref: rack/rack#1952

Similar to `rack.after_reply` but this one is part of the Rack
spec and expect some arguments.

CHANGELOG.md

@@ -1,5 +1,6 @@
 # Unreleased
 
+- Implement `rack.response_finished` (#97).
 - Don't break the `rack.after_reply` callback chain if one callback raises (#97).
 
 # 0.11.1

lib/pitchfork/http_server.rb

@@ -732,7 +732,7 @@ def process_client(client, worker, timeout_handler)
         end
       end
 
-      env["rack.after_reply"] = []
+      env["rack.response_finished"] = env["rack.after_reply"] = []
 
       status, headers, body = @app.call(env)
 
@@ -758,15 +758,19 @@ def process_client(client, worker, timeout_handler)
         client.close # flush and uncork socket immediately, no keepalive
       end
       env
-    rescue => e
-      handle_error(client, e)
+    rescue => application_error
+      handle_error(client, application_error)
       env
     ensure
       if env
-        env["rack.after_reply"].each do |callback|
-          callback.call
-        rescue => error
-          Pitchfork.log_error(@logger, "rack.after_reply error", error)
+        env["rack.response_finished"].each do |callback|
+          if callback.arity == 0
+            callback.call
+          else
+            callback.call(env, status, headers, application_error)
+          end
+        rescue => callback_error
+          Pitchfork.log_error(@logger, "rack.after_reply error", callback_error)
         end
       end
       timeout_handler.finished

test/slow/test_server.rb

@@ -35,20 +35,31 @@ def call(env)
   class TestRackAfterReply
     def initialize
       @called = false
+      @valid_arguments = false
     end
 
     def call(env)
       while env['rack.input'].read(4096)
       end
 
       env["rack.after_reply"] << -> { raise "oops" }
+      env["rack.response_finished"] << method(:check_arguments)
       env["rack.after_reply"] << -> { @called = true }
 
-      [200, { 'content-type' => 'text/plain' }, ["after_reply_called: #{@called}"]]
+      [200, { 'content-type' => 'text/plain' }, ["after_reply_called: #{@called}\nresponse_finished_called: #{@valid_arguments}\n"]]
     rescue Pitchfork::ClientShutdown, Pitchfork::HttpParserError => e
       $stderr.syswrite("#{e.class}: #{e.message} #{e.backtrace.empty?}\n")
       raise e
     end
+
+    def check_arguments(env, status, headers, application_error)
+      return unless env.is_a?(Hash) && env["SERVER_PROTOCOL"]
+      return unless status == 200
+      return unless headers.is_a?(Hash) && headers.key?("content-type")
+      return unless application_error.nil?
+
+      @valid_arguments = true
+    end
   end
 
   def setup
@@ -142,13 +153,15 @@ def test_after_reply
     responses = sock.read(4096)
     assert_match %r{\AHTTP/1.[01] 200\b}, responses
     assert_match %r{^after_reply_called: false}, responses
+    assert_match %r{^response_finished_called: false}, responses
 
     sock = tcp_socket('127.0.0.1', @port)
     sock.syswrite("GET / HTTP/1.0\r\n\r\n")
 
     responses = sock.read(4096)
     assert_match %r{\AHTTP/1.[01] 200\b}, responses
     assert_match %r{^after_reply_called: true}, responses
+    assert_match %r{^response_finished_called: true}, responses
 
     sock.close
   end