From 06979d7d2c67e7d464f584660814dc81fa659e74 Mon Sep 17 00:00:00 2001
From: Zach Anker <zanker@stripe.com>
Date: Mon, 16 Mar 2020 07:30:51 -0700
Subject: [PATCH] Set shutdown error response as a default

---
 lib/puma/configuration.rb | 5 ++++-
 lib/puma/dsl.rb           | 4 ++--
 lib/puma/server.rb        | 9 +--------
 test/test_puma_server.rb  | 2 +-
 4 files changed, 8 insertions(+), 12 deletions(-)

diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb
index 5e01c3b0aa..f2ff990529 100644
--- a/lib/puma/configuration.rb
+++ b/lib/puma/configuration.rb
@@ -13,6 +13,8 @@ module ConfigDefault
     DefaultTCPPort = 9292
     DefaultWorkerTimeout = 60
     DefaultWorkerShutdownTimeout = 30
+
+    DefaultThreadShutdownResponse = [503, {}, ["Request was internally terminated early\n"]].freeze
   end
 
   # A class used for storing "leveled" configuration options.
@@ -186,7 +188,8 @@ def puma_default_options
         :logger => STDOUT,
         :persistent_timeout => Const::PERSISTENT_TIMEOUT,
         :first_data_timeout => Const::FIRST_DATA_TIMEOUT,
-        :raise_exception_on_sigterm => true
+        :raise_exception_on_sigterm => true,
+        :force_shutdown_error_response => DefaultThreadShutdownResponse,
       }
     end
 
diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index 47b737be75..06b0c51aa9 100644
--- a/lib/puma/dsl.rb
+++ b/lib/puma/dsl.rb
@@ -545,8 +545,8 @@ def lowlevel_error_handler(obj=nil, &block)
     # @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)
+      raise "Not a Rack-compatible header (must respond to each)" unless headers.respond_to?(:each)
+      raise "Not a Rack-compatible response (must respond to each)" unless response.respond_to?(:each)
 
       @options[:force_shutdown_error_response] = [
         Integer(status),
diff --git a/lib/puma/server.rb b/lib/puma/server.rb
index f6a9a66dc1..2104f7f7fa 100644
--- a/lib/puma/server.rb
+++ b/lib/puma/server.rb
@@ -593,14 +593,7 @@ def handle_request(req, lines)
           @events.log "Detected force shutdown of a thread, returning 503"
           @events.unknown_error self, e, "Rack app"
 
-          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
-
+          status, headers, res_body = @options[:force_shutdown_error_response]
         rescue Exception => e
           @events.unknown_error self, e, "Rack app", env
 
diff --git a/test/test_puma_server.rb b/test/test_puma_server.rb
index 99461704b7..db1088f66e 100644
--- a/test/test_puma_server.rb
+++ b/test/test_puma_server.rb
@@ -278,7 +278,7 @@ def test_force_shutdown_custom_error_message
   end
 
   def test_force_shutdown_error_default
-    @server = Puma::Server.new @app, @events, {:force_shutdown_after => 2}
+    @server = Puma::Server.new @app, @events, {:force_shutdown_after => 2, :force_shutdown_error_response => Puma::ConfigDefault::DefaultThreadShutdownResponse}
 
     server_run app: ->(env) do
       @server.stop