Merge pull request #98 from Shopify/after-reply-resiliency

Handle errors in `rack.after_reply`

CHANGELOG.md

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

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,11 +758,21 @@ 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
-      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

test/slow/test_server.rb

@@ -35,19 +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
@@ -141,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