diff --git a/lib/puma/binder.rb b/lib/puma/binder.rb index bbcca5bac8..44e923ff7d 100644 --- a/lib/puma/binder.rb +++ b/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 = {} diff --git a/lib/puma/const.rb b/lib/puma/const.rb index 29503f1225..bfe5d69ee7 100644 --- a/lib/puma/const.rb +++ b/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" diff --git a/lib/puma/request.rb b/lib/puma/request.rb index f06744e6db..282f495bee 100644 --- a/lib/puma/request.rb +++ b/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 diff --git a/test/test_rack_server.rb b/test/test_rack_server.rb index c8f3e93514..1848211abf 100644 --- a/test/test_rack_server.rb +++ b/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 }