diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb index 31321ecc0f..cd20281a12 100644 --- a/lib/puma/configuration.rb +++ b/lib/puma/configuration.rb @@ -191,6 +191,8 @@ def initialize(user_options={}, default_options = {}, env = ENV, &block) if block configure(&block) end + + run_mode_hooks end attr_reader :options, :plugins @@ -245,6 +247,8 @@ def puma_options_from_env(env = ENV) def load config_files.each { |config_file| @file_dsl._load_from(config_file) } + run_mode_hooks + @options end @@ -314,6 +318,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 +361,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 @@ -397,11 +418,5 @@ def load_rackup rack_app end - - def self.random_token - require 'securerandom' unless defined?(SecureRandom) - - SecureRandom.hex(16) - end end end diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb index 8601e0217e..58d29226f9 100644 --- a/lib/puma/dsl.rb +++ b/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+ # diff --git a/test/test_config.rb b/test/test_config.rb index 3a20358af6..b40c4a7194 100644 --- a/test/test_config.rb +++ b/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