Add ability to force threads to stop on shutdown. Fixes #938

lib/puma/const.rb

@@ -118,6 +118,11 @@ module Puma
     # sending data back
     WRITE_TIMEOUT = 10
 
+    # How long, after raising the ForceShutdown of a thread during
+    # forced shutdown mode, to wait for the thread to try and finish
+    # up it's work before leaving the thread to die on the vine.
+    SHUTDOWN_GRACE_TIME = 5 # seconds
+
     DATE = "Date".freeze
 
     SCRIPT_NAME = "SCRIPT_NAME".freeze

lib/puma/dsl.rb

@@ -159,6 +159,26 @@ module Puma
       @options[:environment] = environment
     end
 
+    # How long to wait for threads to stop when shutting them
+    # down. Defaults to :forever. Specifying :immediately will cause
+    # Puma to kill the threads immediately.  Otherwise the value
+    # is the number of seconds to wait.
+    #
+    # Puma always waits a few seconds after killing a thread for it to try
+    # to finish up it's work, even in :immediately mode.
+    def force_shutdown_after(val=:forever)
+      i = case val
+          when :forever
+            -1
+          when :immediately
+            0
+          else
+            Integer(val)
+          end
+
+      @options[:force_shutdown_after] = i
+    end
+
     # Code to run before doing a restart. This code should
     # close logfiles, database connections, etc.
     #

lib/puma/server.rb

@@ -571,6 +571,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"]
+
         rescue StandardError => e
           @events.unknown_error self, e, "Rack app"
 
@@ -835,7 +843,13 @@ module Puma
         @events.debug "Drained #{count} additional connections."
       end
 
-      @thread_pool.shutdown if @thread_pool
+      if @thread_pool
+        if timeout = @options[:force_shutdown_after]
+          @thread_pool.shutdown timeout.to_i
+        else
+          @thread_pool.shutdown
+        end
+      end
     end
 
     # Stops the acceptor thread and then causes the worker threads to finish

lib/puma/thread_pool.rb

@@ -5,6 +5,9 @@ module Puma
   #
   class ThreadPool
 
+    class ForceShutdown < RuntimeError
+    end
+
     # Maintain a minimum of +min+ and maximum of +max+ threads
     # in the pool.
     #
@@ -239,7 +242,7 @@ module Puma
 
     # Tell all threads in the pool to exit and wait for them to finish.
     #
-    def shutdown
+    def shutdown(timeout=-1)
       threads = @mutex.synchronize do
         @shutdown = true
         @not_empty.broadcast
@@ -251,7 +254,38 @@ module Puma
         @workers.dup
       end
 
-      threads.each(&:join)
+      case timeout
+      when -1
+        threads.each(&:join)
+      when 0
+        threads.each do |t|
+          t.raise ForceShutdown
+        end
+
+        threads.each do |t|
+          t.join Const::SHUTDOWN_GRACE_TIME
+        end
+      else
+        timeout.times do
+          threads.delete_if do |t|
+            t.join 1
+          end
+
+          if threads.empty?
+            break
+          else
+            sleep 1
+          end
+        end
+
+        threads.each do |t|
+          t.raise ForceShutdown
+        end
+
+        threads.each do |t|
+          t.join Const::SHUTDOWN_GRACE_TIME
+        end
+      end
 
       @spawned = 0
       @workers = []