Add customizable error messages on thread shutdown

History.md

@@ -7,6 +7,7 @@
   * `GC.compact` is called before fork if available (#2093)
   * Add `requests_count` to workers stats. (#2106)
   * Increases maximum URI path length from 2048 to 8196 bytes (#2167)
+  * Added `force_shutdown_error_response` to override the force shutdown response to clients. (#2182)
 
 * Deprecations, Removals and Breaking API Changes
   * `Puma.stats` now returns a Hash instead of a JSON string (#2086)

lib/puma/dsl.rb

@@ -539,6 +539,22 @@ def lowlevel_error_handler(obj=nil, &block)
       @options[:lowlevel_error_handler] = obj
     end
 
+    # When force_shutdown_after is set, this allows you to control the message returned
+    # on a force shutdown to the client.
+    #
+    # @example
+    #   force_shutdown_error_response(500, {"Content-Type" => "application/json"}, [JSON.generate({message: "Server shutdown."})])
+    def force_shutdown_error_response(status, headers, response)
+      raise "Headers must be a hash" unless headers.is_a?(Hash)
+      raise "Response must be enumerable" unless response.is_a?(Enumerable)
+
+      @options[:force_shutdown_error_response] = [
+        Integer(status),
+        headers,
+        response
+      ]
+    end
+
     # This option is used to allow your app and its gems to be
     # properly reloaded when not using preload.
     #

lib/puma/server.rb

@@ -593,9 +593,13 @@ def handle_request(req, lines)
           @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"]
+          if @options[:force_shutdown_error_response]
+            status, headers, res_body = @options[:force_shutdown_error_response]
+          else
+            status = 503
+            headers = {}
+            res_body = ["Request was internally terminated early\n"]
+          end
 
         rescue Exception => e
           @events.unknown_error self, e, "Rack app", env

test/config/app.rb

@@ -7,3 +7,9 @@
 lowlevel_error_handler do |err|
   [200, {}, ["error page"]]
 end
+
+force_shutdown_error_response(
+  500,
+  {"Content-Type" => "application/json"},
+  ["{}"]
+)

test/test_config.rb

@@ -114,6 +114,15 @@ def test_lowlevel_error_handler_DSL
     assert_equal [200, {}, ["error page"]], app.call({})
   end
 
+  def test_force_shutdown_error_response_DSL
+    conf = Puma::Configuration.new do |c|
+      c.load "test/config/app.rb"
+    end
+    conf.load
+
+    assert_equal [500, {"Content-Type" => "application/json"}, ["{}"]], conf.options[:force_shutdown_error_response]
+  end
+
   def test_allow_users_to_override_default_options
     conf = Puma::Configuration.new(restart_cmd: 'bin/rails server')
 

test/test_puma_server.rb

@@ -261,6 +261,36 @@ def test_doesnt_print_backtrace_in_production
     assert_match(/HTTP\/1.0 500 Internal Server Error/, data)
   end
 
+  def test_force_shutdown_custom_error_message
+    shutdown_response = [500, {"Content-Type" => "application/json"}, ["{}\n"]]
+    @server = Puma::Server.new @app, @events, {:force_shutdown_error_response => shutdown_response, :force_shutdown_after => 2}
+
+    server_run app: ->(env) do
+      @server.stop
+      sleep 10
+    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 10
+    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(/Request was internally terminated early\n$/, 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}