Add single and cluster configuration hooks

lib/puma/configuration.rb

@@ -191,6 +191,10 @@ def initialize(user_options={}, default_options = {}, env = ENV, &block)
       if block
         configure(&block)
       end
+
+      run_mode_hooks
+
+      set_conditional_default_options
     end
 
     attr_reader :options, :plugins
@@ -245,6 +249,10 @@ def puma_options_from_env(env = ENV)
     def load
       config_files.each { |config_file| @file_dsl._load_from(config_file) }
 
+      run_mode_hooks
+
+      set_conditional_default_options
+
       @options
     end
 
@@ -314,6 +322,17 @@ def load_plugin(name)
       @plugins.create name
     end
 
+    def run_mode_hooks
+      workers_before = @options[:workers]
+      key = workers_before > 0 ? :cluster : :single
+
+      @options.all_of(key).each { |block| configure(&block) }
+
+      unless workers_before == @options[:workers]
+        raise ArgumentError, "cannot change the number of workers inside a #{key} hook"
+      end
+    end
+
     # @param key [:Symbol] hook to run
     # @param arg [Launcher, Int] `:on_restart` passes Launcher
     #
@@ -346,6 +365,12 @@ def self.temp_path
       "#{Dir.tmpdir}/puma-status-#{t}-#{$$}"
     end
 
+    def self.random_token
+      require 'securerandom' unless defined?(SecureRandom)
+
+      SecureRandom.hex(16)
+    end
+
     private
 
     def require_processor_counter
@@ -398,10 +423,9 @@ def load_rackup
       rack_app
     end
 
-    def self.random_token
-      require 'securerandom' unless defined?(SecureRandom)
-
-      SecureRandom.hex(16)
+    def set_conditional_default_options
+      @options.default_options[:preload_app] = !@options[:prune_bundler] &&
+        (@options[:workers] > 1) && Puma.forkable?
     end
   end
 end

lib/puma/dsl.rb

@@ -712,6 +712,24 @@ def silence_fork_callback_warning
       @options[:silence_fork_callback_warning] = true
     end
 
+    # Code to run after loading the global configuration.
+    #
+    # @note Single mode only.
+    #
+    def single(&block)
+      @options[:single] ||= []
+      @options[:single] << block
+    end
+
+    # Code to run after loading the global configuration.
+    #
+    # @note Cluster mode only.
+    #
+    def cluster(&block)
+      @options[:cluster] ||= []
+      @options[:cluster] << block
+    end
+
     # Code to run immediately before master process
     # forks workers (once on boot). These hooks can block if necessary
     # to wait for background operations unknown to Puma to finish before
@@ -991,6 +1009,7 @@ def lowlevel_error_handler(obj=nil, &block)
     # new Bundler context and thus can float around as the release
     # dictates.
     #
+    # @note Cluster mode only.
     # @note This is incompatible with +preload_app!+.
     # @note This is only supported for RubyGems 2.2+
     #

test/test_config.rb

@@ -468,6 +468,88 @@ def test_config_raise_exception_on_sigterm
     assert_equal conf.options[:raise_exception_on_sigterm], true
   end
 
+  def test_run_mode_hooks_with_unspecified_workers
+    evals = []
+    Puma::Configuration.new do |c|
+      c.single { evals << :single }
+      c.cluster { evals << :cluster }
+    end
+    assert_equal [:single], evals
+  end
+
+  def test_run_mode_hooks_with_zero_workers
+    evals = []
+    Puma::Configuration.new do |c|
+      c.workers 0
+      c.single { evals << :single }
+      c.cluster { evals << :cluster }
+    end
+    assert_equal [:single], evals
+  end
+
+  def test_run_mode_hooks_with_positive_workers
+    evals = []
+    Puma::Configuration.new do |c|
+      c.workers 2
+      c.single { evals << :single }
+      c.cluster { evals << :cluster }
+    end
+    assert_equal [:cluster], evals
+  end
+
+  def test_run_mode_hooks_with_positive_workers_after_hooks
+    evals = []
+    Puma::Configuration.new do |c|
+      c.single { evals << :single }
+      c.cluster { evals << :cluster }
+      c.workers 2
+    end
+    assert_equal [:cluster], evals
+  end
+
+  def test_run_mode_hooks_with_workers_overwritten
+    evals = []
+    Puma::Configuration.new do |c|
+      c.workers 2
+      c.single { evals << :single }
+      c.cluster { evals << :cluster }
+      c.workers 0
+    end
+    assert_equal [:single], evals
+  end
+
+  def test_run_mode_hooks_with_workers_overwritten_in_hook
+    error = assert_raises ArgumentError do
+      Puma::Configuration.new do |c|
+        c.workers 2
+        c.cluster { c.workers 0 }
+      end
+    end
+    assert_equal "cannot change the number of workers inside a cluster hook", error.message
+  end
+
+  def test_run_mode_hooks_with_workers_duplicated_in_hook
+    evals = []
+    Puma::Configuration.new do |c|
+      c.workers 2
+      c.single { evals << :single }
+      c.cluster { c.workers 2; evals << :cluster }
+    end
+    assert_equal [:cluster], evals
+  end
+
+  def test_run_mode_hooks_with_multiple_hooks
+    evals = []
+    Puma::Configuration.new do |c|
+      c.workers 2
+      c.single { evals << :single_1 }
+      c.cluster { evals << :cluster_1 }
+      c.single { evals << :single_2 }
+      c.cluster { evals << :cluster_2 }
+    end
+    assert_equal [:cluster_1, :cluster_2], evals
+  end
+
   def test_run_hooks_on_restart_hook
     assert_run_hooks :on_restart
   end