Handle thread shutdown responses via low level error messages (#2203)

* Handle thread shutdown responses via low level error messages

* Split ForceShutdown and Exception handling

* Tighten sleep to 5 seconds

History.md

@@ -9,6 +9,7 @@
   * Add `requests_count` to workers stats. (#2106)
   * Increases maximum URI path length from 2048 to 8196 bytes (#2167)
   * Sending SIGWINCH to any Puma worker now prints currently active threads and their backtraces (#2195)
+  * Force shutdown responses can be overridden by using the `lowlevel_error_handler` config (#2203)
 
 * Deprecations, Removals and Breaking API Changes
   * `Puma.stats` now returns a Hash instead of a JSON string (#2086)

lib/puma/server.rb

@@ -590,17 +590,14 @@ module Puma
             return :async
           end
         rescue ThreadPool::ForceShutdown => e
-          @events.log "Detected force shutdown of a thread, returning 503"
-          @events.unknown_error self, e, "Rack app"
-
-          status = 503
-          headers = {}
-          res_body = ["Request was internally terminated early\n"]
+          @events.unknown_error self, e, "Rack app", env
+          @events.log "Detected force shutdown of a thread"
 
+          status, headers, res_body = lowlevel_error(e, env, 503)
         rescue Exception => e
           @events.unknown_error self, e, "Rack app", env
 
-          status, headers, res_body = lowlevel_error(e, env)
+          status, headers, res_body = lowlevel_error(e, env, 500)
         end
 
         content_length = nil
@@ -807,19 +804,21 @@ module Puma
 
     # A fallback rack response if +@app+ raises as exception.
     #
-    def lowlevel_error(e, env)
+    def lowlevel_error(e, env, status=500)
       if handler = @options[:lowlevel_error_handler]
         if handler.arity == 1
           return handler.call(e)
-        else
+        elsif handler.arity == 2
           return handler.call(e, env)
+        else
+          return handler.call(e, env, status)
         end
       end
 
       if @leak_stack_on_error
-        [500, {}, ["Puma caught this error: #{e.message} (#{e.class})\n#{e.backtrace.join("\n")}"]]
+        [status, {}, ["Puma caught this error: #{e.message} (#{e.class})\n#{e.backtrace.join("\n")}"]]
       else
-        [500, {}, ["An unhandled lowlevel error occurred. The application logs may have details.\n"]]
+        [status, {}, ["An unhandled lowlevel error occurred. The application logs may have details.\n"]]
       end
     end
 

test/test_puma_server.rb

@@ -261,6 +261,36 @@ EOF
     assert_match(/HTTP\/1.0 500 Internal Server Error/, data)
   end
 
+  def test_force_shutdown_custom_error_message
+    handler = lambda {|err, env, status| [500, {"Content-Type" => "application/json"}, ["{}\n"]]}
+    @server = Puma::Server.new @app, @events, {:lowlevel_error_handler => handler, :force_shutdown_after => 2}
+
+    server_run app: ->(env) do
+      @server.stop
+      sleep 5
+    end
+
+    data = send_http_and_read "GET / HTTP/1.0\r\n\r\n"
+
+    assert_match(/HTTP\/1.0 500 Internal Server Error/, data)
+    assert_match(/Content-Type: application\/json/, data)
+    assert_match(/{}\n$/, data)
+  end
+
+  def test_force_shutdown_error_default
+    @server = Puma::Server.new @app, @events, {:force_shutdown_after => 2}
+
+    server_run app: ->(env) do
+      @server.stop
+      sleep 5
+    end
+
+    data = send_http_and_read "GET / HTTP/1.0\r\n\r\n"
+
+    assert_match(/HTTP\/1.0 503 Service Unavailable/, data)
+    assert_match(/Puma caught this error.+Puma::ThreadPool::ForceShutdown/, data)
+  end
+
   def test_prints_custom_error
     re = lambda { |err| [302, {'Content-Type' => 'text', 'Location' => 'foo.html'}, ['302 found']] }
     @server = Puma::Server.new @app, @events, {:lowlevel_error_handler => re}
@@ -287,6 +317,21 @@ EOF
     assert_match(/HTTP\/1.0 302 Found/, data)
   end
 
+  def test_leh_has_status
+    re = lambda { |err, env, status|
+      raise "Cannot find status" unless status
+      [302, {'Content-Type' => 'text', 'Location' => 'foo.html'}, ['302 found']]
+    }
+
+    @server = Puma::Server.new @app, @events, {:lowlevel_error_handler => re}
+
+    server_run app: ->(env) { raise "don't leak me bro" }
+
+    data = send_http_and_read "GET / HTTP/1.0\r\n\r\n"
+
+    assert_match(/HTTP\/1.0 302 Found/, data)
+  end
+
   def test_custom_http_codes_10
     server_run app: ->(env) { [449, {}, [""]] }