From e5150f07927e96aa5aae8a65db116472de1d9d63 Mon Sep 17 00:00:00 2001
From: Max Gulyaev <maxim@yotpo.com>
Date: Wed, 26 Feb 2014 16:06:58 +0200
Subject: [PATCH] added after worker boot hook

---
 examples/config.rb        |  9 +++++++++
 lib/puma/cli.rb           |  5 +++--
 lib/puma/cluster.rb       |  3 ++-
 lib/puma/configuration.rb | 14 ++++++++++++--
 4 files changed, 26 insertions(+), 5 deletions(-)

diff --git a/examples/config.rb b/examples/config.rb
index 251dce5b..231ab54e 100644
--- a/examples/config.rb
+++ b/examples/config.rb
@@ -113,6 +113,15 @@
 #   puts 'On worker boot...'
 # end
 
+# Code to run when a worker boots to setup the process after booting
+# the app.
+#
+# This can be called multiple times to add hooks.
+#
+# after_worker_boot do
+#   puts 'On worker boot...'
+# end
+
 # === Puma control rack application ===
 
 # Start the puma control rack application on “url”. This application can
diff --git a/lib/puma/cli.rb b/lib/puma/cli.rb
index 475ec2e8..29d71b75 100644
--- a/lib/puma/cli.rb
+++ b/lib/puma/cli.rb
@@ -99,7 +99,8 @@ module Puma
         :binds => [],
         :workers => 0,
         :daemon => false,
-        :worker_boot => []
+        :before_worker_boot => [],
+        :after_worker_boot => []
       }
 
       @parser = OptionParser.new do |o|
@@ -220,7 +221,7 @@ module Puma
 
         cfg = @config.dup
 
-        [ :logger, :worker_boot, :on_restart ].each { |o| cfg.options.delete o }
+        [ :logger, :before_worker_boot, :after_worker_boot, :on_restart ].each { |o| cfg.options.delete o }
 
         state["config"] = cfg
 
diff --git a/lib/puma/cluster.rb b/lib/puma/cluster.rb
index bd725b68..f7ba9bd9 100644
--- a/lib/puma/cluster.rb
+++ b/lib/puma/cluster.rb
@@ -90,6 +90,7 @@ module Puma
         pid = fork { worker(idx, upgrade, master) }
         @cli.debug "Spawned worker: #{pid}"
         @workers << Worker.new(idx, pid, @phase)
+        @options[:after_worker_boot].each { |h| h.call }
       end
 
       if diff > 0
@@ -186,7 +187,7 @@ module Puma
 
       # Invoke any worker boot hooks so they can get
       # things in shape before booting the app.
-      hooks = @options[:worker_boot]
+      hooks = @options[:before_worker_boot]
       hooks.each { |h| h.call(index) }
 
       server = start_server
diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb
index 23249111..b4ad8160 100644
--- a/lib/puma/configuration.rb
+++ b/lib/puma/configuration.rb
@@ -21,7 +21,8 @@ module Puma
       @options[:mode] ||= :http
       @options[:binds] ||= []
       @options[:on_restart] ||= []
-      @options[:worker_boot] ||= []
+      @options[:before_worker_boot] ||= []
+      @options[:after_worker_boot] ||= []
       @options[:worker_timeout] ||= DefaultWorkerTimeout
     end
 
@@ -306,7 +307,16 @@ module Puma
       # This can be called multiple times to add hooks.
       #
       def on_worker_boot(&block)
-        @options[:worker_boot] << block
+        @options[:before_worker_boot] << block
+      end
+
+      # *Cluster mode only* Code to run when a worker boots to setup
+      # the process after booting the app.
+      #
+      # This can be called multiple times to add hooks.
+      #
+      def after_worker_boot(&block)
+        @options[:after_worker_boot] << block
       end
 
       # The directory to operate out of.