lib/pitchfork/http_server.rb in pitchfork-0.11.1 vs lib/pitchfork/http_server.rb in pitchfork-0.12.0

@@ -730,11 +730,11 @@
         env["rack.early_hints"] = lambda do |headers|
           e103_response_write(client, headers)
         end
       end
 
-      env["rack.after_reply"] = []
+      env["rack.response_finished"] = env["rack.after_reply"] = []
 
       status, headers, body = @app.call(env)
 
       begin
         return env if @request.hijacked?
@@ -756,14 +756,24 @@
         rescue Errno::ENOTCONN
         end
         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
-      env["rack.after_reply"].each(&:call) if env
+      if env
+        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
       env
     end
 
     def nuke_listeners!(readers)