Merge pull request #566 from sheltond/master

Added on_worker_shutdown mechanism

lib/puma/cli.rb

@@ -99,6 +99,7 @@ module Puma
         :binds => [],
         :workers => 0,
         :daemon => false,
+        :before_worker_shutdown => [],
         :before_worker_boot => [],
         :after_worker_boot => []
       }
@@ -228,7 +229,7 @@ module Puma
 
         cfg = @config.dup
 
-        [ :logger, :before_worker_boot, :after_worker_boot, :on_restart ].each { |o| cfg.options.delete o }
+        [ :logger, :before_worker_shutdown, :before_worker_boot, :after_worker_boot, :on_restart ].each { |o| cfg.options.delete o }
 
         state["config"] = cfg
 

lib/puma/cluster.rb

@@ -73,7 +73,7 @@ module Puma
 
       def term
         begin
-          if @first_term_sent && (Time.new - @first_term_sent) > 30
+          if @first_term_sent && (Time.new - @first_term_sent) > @options[:worker_shutdown_timeout]
             @signal = "KILL"
           else
             @first_term_sent ||= Time.new
@@ -228,6 +228,10 @@ module Puma
 
       server.run.join
 
+      # Invoke any worker shutdown hooks so they can prevent the worker
+      # exiting until any background operations are completed
+      hooks = @options[:before_worker_shutdown]
+      hooks.each { |h| h.call(index) }
     ensure
       @worker_write.close
     end

lib/puma/configuration.rb

@@ -15,15 +15,18 @@ module Puma
     DefaultTCPHost = "0.0.0.0"
     DefaultTCPPort = 9292
     DefaultWorkerTimeout = 60
+    DefaultWorkerShutdownTimeout = 30
 
     def initialize(options)
       @options = options
       @options[:mode] ||= :http
       @options[:binds] ||= []
       @options[:on_restart] ||= []
+      @options[:before_worker_shutdown] ||= []
       @options[:before_worker_boot] ||= []
       @options[:after_worker_boot] ||= []
       @options[:worker_timeout] ||= DefaultWorkerTimeout
+      @options[:worker_shutdown_timeout] ||= DefaultWorkerShutdownTimeout
     end
 
     attr_reader :options
@@ -305,6 +308,17 @@ module Puma
         @options[:workers] = count.to_i
       end
 
+      # *Cluster mode only* Code to run immediately before a worker shuts
+      # down (after it has finished processing HTTP requests). These hooks
+      # can block if necessary to wait for background operations unknown
+      # to puma to finish before the process terminates.
+      #
+      # This can be called multiple times to add hooks.
+      #
+      def on_worker_shutdown(&block)
+        @options[:before_worker_shutdown] << block
+      end
+
       # *Cluster mode only* Code to run when a worker boots to setup
       # the process before booting the app.
       #
@@ -376,6 +390,11 @@ module Puma
       def worker_timeout(timeout)
         @options[:worker_timeout] = timeout
       end
+
+      # *Cluster mode only* Set the timeout for worker shutdown
+      def worker_shutdown_timeout(timeout)
+        @options[:worker_shutdown_timeout] = timeout
+      end
     end
   end
 end

lib/puma/thread_pool.rb

@@ -180,7 +180,9 @@ module Puma
 
       # Use this instead of #each so that we don't stop in the middle
       # of each and see a mutated object mid #each
-      @workers.first.join until @workers.empty?
+      if !@workers.empty?
+          @workers.first.join until @workers.empty?
+      end
 
       @spawned = 0
       @workers = []