Allow run_hooks to pass a hash to blocks for use later (#2917)

* Update dsl.rb

* Update worker.rb

* Update configuration.rb

* Add test_hook_data

lib/puma/cluster/worker.rb

@@ -23,6 +23,7 @@ def initialize(index:, master:, launcher:, pipes:, server: nil)
         @fork_pipe = pipes[:fork_pipe]
         @wakeup = pipes[:wakeup]
         @server = server
+        @hook_data = {}
       end
 
       def run
@@ -52,7 +53,7 @@ def run
 
         # Invoke any worker boot hooks so they can get
         # things in shape before booting the app.
-        @launcher.config.run_hooks(:before_worker_boot, index, @launcher.log_writer)
+        @launcher.config.run_hooks(:before_worker_boot, index, @launcher.log_writer, @hook_data)
 
         begin
         server = @server ||= start_server
@@ -84,7 +85,7 @@ def run
                 if restart_server.length > 0
                   restart_server.clear
                   server.begin_restart(true)
-                  @launcher.config.run_hooks(:before_refork, nil, @launcher.log_writer)
+                  @launcher.config.run_hooks(:before_refork, nil, @launcher.log_writer, @hook_data)
                 end
               elsif idx == 0 # restart server
                 restart_server << true << false
@@ -138,7 +139,7 @@ def run
 
         # Invoke any worker shutdown hooks so they can prevent the worker
         # exiting until any background operations are completed
-        @launcher.config.run_hooks(:before_worker_shutdown, index, @launcher.log_writer)
+        @launcher.config.run_hooks(:before_worker_shutdown, index, @launcher.log_writer, @hook_data)
       ensure
         @worker_write << "t#{Process.pid}\n" rescue nil
         @worker_write.close
@@ -147,7 +148,7 @@ def run
       private
 
       def spawn_worker(idx)
-        @launcher.config.run_hooks(:before_worker_fork, idx, @launcher.log_writer)
+        @launcher.config.run_hooks(:before_worker_fork, idx, @launcher.log_writer, @hook_data)
 
         pid = fork do
           new_worker = Worker.new index: idx,
@@ -165,7 +166,7 @@ def spawn_worker(idx)
           exit! 1
         end
 
-        @launcher.config.run_hooks(:after_worker_fork, idx, @launcher.log_writer)
+        @launcher.config.run_hooks(:after_worker_fork, idx, @launcher.log_writer, @hook_data)
         pid
       end
     end

lib/puma/configuration.rb

@@ -302,10 +302,18 @@ def load_plugin(name)
       @plugins.create name
     end
 
-    def run_hooks(key, arg, log_writer)
+    # @param key [:Symbol] hook to run
+    # @param arg [Launcher, Int] `:on_restart` passes Launcher
+    #
+    def run_hooks(key, arg, log_writer, hook_data = nil)
       @options.all_of(key).each do |b|
         begin
-          b.call arg
+          if Array === b
+            hook_data[b[1]] ||= Hash.new
+            b[0].call arg, hook_data[b[1]]
+          else
+            b.call arg
+          end
         rescue => e
           log_writer.log "WARNING hook #{key} failed with exception (#{e.class}) #{e.message}"
           log_writer.debug e.backtrace.join("\n")

lib/puma/dsl.rb

@@ -32,7 +32,25 @@ module Puma
   # You can also find many examples being used by the test suite in
   # +test/config+.
   #
+  # Puma v6 adds the option to specify a key name (String or Symbol) to the
+  # hooks that run inside the forked workers.  All the hooks run inside the
+  # {Puma::Cluster::Worker#run} method.
+  #
+  # Previously, the worker index and the LogWriter instance were passed to the
+  # hook blocks/procs.  If a key name is specified, a hash is passed as the last
+  # parameter.  This allows storage of data, typically objects that are created
+  # before the worker that need to be passed to the hook when the worker is shutdown.
+  #
+  # The following hooks have been updated:
+  #
+  #     | DSL Method         |  Options Key            | Fork Block Location |
+  #     | on_worker_boot     | :before_worker_boot     | inside, before      |
+  #     | on_worker_shutdown | :before_worker_shutdown | inside, after       |
+  #     | on_refork          | :before_refork          | inside              |
+  #
   class DSL
+    ON_WORKER_KEY = [String, Symbol].freeze
+
     # convenience method so logic can be used in CI
     # @see ssl_bind
     #
@@ -499,7 +517,7 @@ def threads(min, max)
     #   ssl_bind '127.0.0.1', '9292', {
     #     cert_pem: File.read(path_to_cert),
     #     key_pem: File.read(path_to_key),
-    #     reuse: {size: 2_000, timeout: 20}
+    #     reuse: {size: 2_000, timeout: 20} # optional
     #   }
     #
     # @example For JRuby, two keys are required: +keystore+ & +keystore_pass+
@@ -592,9 +610,8 @@ def before_fork(&block)
     #   on_worker_boot do
     #     puts 'Before worker boot...'
     #   end
-    def on_worker_boot(&block)
-      @options[:before_worker_boot] ||= []
-      @options[:before_worker_boot] << block
+    def on_worker_boot(key = nil, &block)
+      process_hook :before_worker_boot, key, block, 'on_worker_boot'
     end
 
     # Code to run immediately before a worker shuts
@@ -609,9 +626,8 @@ def on_worker_boot(&block)
     #   on_worker_shutdown do
     #     puts 'On worker shutdown...'
     #   end
-    def on_worker_shutdown(&block)
-      @options[:before_worker_shutdown] ||= []
-      @options[:before_worker_shutdown] << block
+    def on_worker_shutdown(key = nil, &block)
+      process_hook :before_worker_shutdown, key, block, 'on_worker_shutdown'
     end
 
     # Code to run in the master right before a worker is started. The worker's
@@ -625,8 +641,7 @@ def on_worker_shutdown(&block)
     #     puts 'Before worker fork...'
     #   end
     def on_worker_fork(&block)
-      @options[:before_worker_fork] ||= []
-      @options[:before_worker_fork] << block
+      process_hook :before_worker_fork, nil, block, 'on_worker_fork'
     end
 
     # Code to run in the master after a worker has been started. The worker's
@@ -640,8 +655,7 @@ def on_worker_fork(&block)
     #     puts 'After worker fork...'
     #   end
     def after_worker_fork(&block)
-      @options[:after_worker_fork] ||= []
-      @options[:after_worker_fork] << block
+      process_hook :after_worker_fork, nil, block, 'after_worker_fork'
     end
 
     alias_method :after_worker_boot, :after_worker_fork
@@ -664,9 +678,8 @@ def after_worker_fork(&block)
     #   end
     # @version 5.0.0
     #
-    def on_refork(&block)
-      @options[:before_refork] ||= []
-      @options[:before_refork] << block
+    def on_refork(key = nil, &block)
+      process_hook :before_refork, key, block, 'on_refork'
     end
 
     # Code to run out-of-band when the worker is idle.
@@ -679,8 +692,7 @@ def on_refork(&block)
     #
     # This can be called multiple times to add several hooks.
     def out_of_band(&block)
-      @options[:out_of_band] ||= []
-      @options[:out_of_band] << block
+      process_hook :out_of_band, nil, block, 'out_of_band'
     end
 
     # The directory to operate out of.
@@ -1026,5 +1038,16 @@ def add_pem_values_to_options_store(opts)
         end
       end
     end
+
+    def process_hook(options_key, key, block, meth)
+      @options[options_key] ||= []
+      if ON_WORKER_KEY.include? key.class
+        @options[options_key] << [block, key.to_sym]
+      elsif key.nil?
+        @options[options_key] << block
+      else
+        raise "'#{method}' key must be String or Symbol"
+      end
+    end
   end
 end

test/config/hook_data.rb

@@ -0,0 +1,9 @@
+workers 2
+
+on_worker_boot(:test) do |index, data|
+  data[:test] = index
+end
+
+on_worker_shutdown(:test) do |index, data|
+  File.write "hook_data-#{index}.txt", "index #{index} data #{data[:test]}", mode: 'wb:UTF-8'
+end

test/helpers/integration.rb

@@ -67,7 +67,7 @@ def silent_and_checked_system_command(*args)
     assert(system(*args, out: File::NULL, err: File::NULL))
   end
 
-  def cli_server(argv, unix: false, config: nil, merge_err: false)
+  def cli_server(argv, unix: false, config: nil, merge_err: false, log: false)
     if config
       config_file = Tempfile.new(%w(config .rb))
       config_file.write config
@@ -86,7 +86,7 @@ def cli_server(argv, unix: false, config: nil, merge_err: false)
     else
       @server = IO.popen(cmd, "r")
     end
-    wait_for_server_to_boot
+    wait_for_server_to_boot(log: log)
     @pid = @server.pid
     @server
   end

test/test_integration_cluster.rb

@@ -421,6 +421,22 @@ def test_culling_strategy_oldest_fork_worker
     assert_match(/Worker 1 \(PID: \d+\) terminating/, line)
   end
 
+  def test_hook_data
+    skip_unless_signal_exist? :TERM
+
+    cli_server "-C test/config/hook_data.rb test/rackup/hello.ru"
+    get_worker_pids 0, 2
+    stop_server
+
+    file = 'hook_data-0.txt'
+    assert_equal 'index 0 data 0', File.read(file, mode: 'rb:UTF-8')
+    File.unlink file if File.file? file
+
+    file = 'hook_data-1.txt'
+    assert_equal 'index 1 data 1', File.read(file, mode: 'rb:UTF-8')
+    File.unlink file if File.file? file
+  end
+
   private
 
   def worker_timeout(timeout, iterations, details, config)