Fix ThreadPool#shutdown timeout accuracy

Gemfile

@@ -10,6 +10,7 @@ gem "rack", "~> 1.6"
 gem "minitest", "~> 5.11"
 gem "minitest-retry"
 gem "minitest-proveit"
+gem "minitest-stub-const"
 
 gem "jruby-openssl", :platform => "jruby"
 

History.md

@@ -32,6 +32,7 @@
   * Pass queued requests to thread pool on server shutdown (#2122)
   * Fixed a few minor concurrency bugs in ThreadPool that may have affected non-GVL Rubies (#2220)
   * Fix `out_of_band` hook never executed if the number of worker threads is > 1 (#2177)
+  * Fix ThreadPool#shutdown timeout accuracy (#2221)
 
 * Refactor
   * Remove unused loader argument from Plugin initializer (#2095)

lib/puma/dsl.rb

@@ -243,7 +243,7 @@ def force_shutdown_after(val=:forever)
           when :immediately
             0
           else
-            Integer(val)
+            Float(val)
           end
 
       @options[:force_shutdown_after] = i

lib/puma/thread_pool.rb

@@ -277,6 +277,9 @@ def auto_reap!(timeout=5)
     end
 
     # Tell all threads in the pool to exit and wait for them to finish.
+    # Wait +timeout+ seconds then raise +ForceShutdown+ in remaining threads.
+    # Next, wait an extra +grace+ seconds then force-kill remaining threads.
+    # Finally, wait +kill_grace+ seconds for remaining threads to exit.
     #
     def shutdown(timeout=-1)
       threads = with_mutex do
@@ -295,27 +298,26 @@ def shutdown(timeout=-1)
         # Wait for threads to finish without force shutdown.
         threads.each(&:join)
       else
-        # Wait for threads to finish after n attempts (+timeout+).
-        # If threads are still running, it will forcefully kill them.
-        timeout.times do
-          threads.delete_if do |t|
-            t.join 1
-          end
-
-          if threads.empty?
-            break
-          else
-            sleep 1
+        join = ->(timeout) do
+          start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+          threads.reject! do |t|
+            elapsed = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
+            t.join timeout - elapsed
           end
         end
 
+        # Wait +timeout+ seconds for threads to finish.
+        join.call(timeout)
+
+        # If threads are still running, raise ForceShutdown and wait to finish.
         threads.each do |t|
           t.raise ForceShutdown
         end
+        join.call(SHUTDOWN_GRACE_TIME)
 
-        threads.each do |t|
-          t.join SHUTDOWN_GRACE_TIME
-        end
+        # If threads are _still_ running, forcefully kill them and wait to finish.
+        threads.each(&:kill)
+        join.call(1)
       end
 
       @spawned = 0

test/helper.rb

@@ -14,6 +14,7 @@
 require "minitest/autorun"
 require "minitest/pride"
 require "minitest/proveit"
+require "minitest/stub_const"
 require_relative "helpers/apps"
 
 Thread.abort_on_exception = true

test/test_puma_server.rb

@@ -934,16 +934,16 @@ def test_shutdown_data_timeout
 
   # Requests still pending after `force_shutdown_after` should have connection closed (408 w/pending POST body).
   def test_force_shutdown
-    shutdown_requests request_delay: 4, response: nil, force_shutdown_after: 1
-    shutdown_requests request_delay: 4, response: nil, force_shutdown_after: 1, queue_requests: false
-    shutdown_requests request_delay: 4, response: /408/, force_shutdown_after: 1, post: true
+    shutdown_requests request_delay: 4, response: nil, force_shutdown_after: 3
+    shutdown_requests request_delay: 4, response: nil, force_shutdown_after: 3, queue_requests: false
+    shutdown_requests request_delay: 4, response: /408/, force_shutdown_after: 3, post: true
   end
 
   # App-responses still pending during `force_shutdown_after` should return 503
   # (uncaught Puma::ThreadPool::ForceShutdown exception).
   def test_force_shutdown_app
-    shutdown_requests app_delay: 3, response: /503/, force_shutdown_after: 1
-    shutdown_requests app_delay: 3, response: /503/, force_shutdown_after: 1, queue_requests: false
+    shutdown_requests app_delay: 3, response: /503/, force_shutdown_after: 3
+    shutdown_requests app_delay: 3, response: /503/, force_shutdown_after: 3, queue_requests: false
   end
 
   def test_http11_connection_header_queue

test/test_thread_pool.rb

@@ -240,4 +240,30 @@ def test_waiting_on_startup
     pool = new_pool(1, 2)
     assert_equal 1, pool.waiting
   end
+
+  def test_shutdown_with_grace
+    timeout = 0.01
+    grace = 0.01
+
+    rescued = []
+    pool = mutex_pool(2, 2) do
+      begin
+        pool.signal
+        sleep
+      rescue Puma::ThreadPool::ForceShutdown
+        rescued << Thread.current
+        sleep
+      end
+    end
+
+    pool << 1
+    pool << 2
+
+    Puma::ThreadPool.stub_const(:SHUTDOWN_GRACE_TIME, grace) do
+      pool.shutdown(timeout)
+    end
+    assert_equal 0, pool.spawned
+    assert_equal 2, rescued.length
+    refute rescued.any?(&:alive?)
+  end
 end