Set the worker process count automatically with WEB_CONCURRENCY=auto (#…

…3439) Requires the concurrent-ruby gem, displays an error message if it is not available and WEB_CONCURRENCY is set to "auto".
puma · Nov 4, 2024 · 6f26742 · 6f26742
1 parent 5ac3f40
commit 6f26742
Show file tree

Hide file tree

Showing 7 changed files with 66 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -118,6 +118,8 @@ $ WEB_CONCURRENCY=3 puma -t 8:32
 
 Note that threads are still used in clustered mode, and the `-t` thread flag setting is per worker, so `-w 2 -t 16:16` will spawn 32 threads in total, with 16 in each worker process.
 
+If the `WEB_CONCURRENCY` environment variable is set to `"auto"` and the `concurrent-ruby` gem is available in your application, Puma will set the worker process count to the result of [available processors](https://ruby-concurrency.github.io/concurrent-ruby/master/Concurrent.html#available_processor_count-class_method).
+
 For an in-depth discussion of the tradeoffs of thread and process count settings, [see our docs](https://github.com/puma/puma/blob/9282a8efa5a0c48e39c60d22ca70051a25df9f55/docs/kubernetes.md#workers-per-pod-and-other-config-issues).
 
 In clustered mode, Puma can "preload" your application. This loads all the application code *prior* to forking. Preloading reduces total memory usage of your application via an operating system feature called [copy-on-write](https://en.wikipedia.org/wiki/Copy-on-write).

diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb
@@ -227,7 +227,12 @@ def puma_default_options(env = ENV)
     def puma_options_from_env(env = ENV)
       min = env['PUMA_MIN_THREADS'] || env['MIN_THREADS']
       max = env['PUMA_MAX_THREADS'] || env['MAX_THREADS']
-      workers = env['WEB_CONCURRENCY']
+      workers = if env['WEB_CONCURRENCY'] == 'auto'
+        require_processor_counter
+        ::Concurrent.available_processor_count
+      else
+        env['WEB_CONCURRENCY']
+      end
 
       {
         min_threads: min && Integer(min),
@@ -341,6 +346,16 @@ def self.temp_path
 
     private
 
+    def require_processor_counter
+      require 'concurrent/utility/processor_counter'
+    rescue LoadError
+      warn <<~MESSAGE
+        WEB_CONCURRENCY=auto requires the "concurrent-ruby" gem to be installed.
+        Please add "concurrent-ruby" to your Gemfile.
+      MESSAGE
+      raise
+    end
+
     # Load and use the normal Rack builder if we can, otherwise
     # fallback to our minimal version.
     def rack_builder

diff --git a/test/helpers/integration.rb b/test/helpers/integration.rb
@@ -76,6 +76,15 @@ def silent_and_checked_system_command(*args)
     assert(system(*args, out: File::NULL, err: File::NULL))
   end
 
+  def with_unbundled_env
+    bundler_ver = Gem::Version.new(Bundler::VERSION)
+    if bundler_ver < Gem::Version.new('2.1.0')
+      Bundler.with_clean_env { yield }
+    else
+      Bundler.with_unbundled_env { yield }
+    end
+  end
+
   def cli_server(argv,  # rubocop:disable Metrics/ParameterLists
       unix: false,      # uses a UNIXSocket for the server listener when true
       config: nil,      # string to use for config file

diff --git a/test/test_web_concurrency_auto.rb b/test/test_web_concurrency_auto.rb
@@ -0,0 +1,29 @@
+require_relative "helper"
+require_relative "helpers/integration"
+
+class TestWebConcurrencyAuto < TestIntegration
+
+  def teardown
+    return if skipped?
+    super
+  end
+
+  def test_web_concurrency_with_concurrent_ruby_available
+    skip_unless :fork
+    env = {
+      "BUNDLE_GEMFILE" => "#{__dir__}/web_concurrency_test/Gemfile",
+      "WEB_CONCURRENCY" => "auto"
+    }
+    Dir.chdir("#{__dir__}/web_concurrency_test") do
+      with_unbundled_env do
+        silent_and_checked_system_command("bundle config --local path vendor/bundle")
+        silent_and_checked_system_command("bundle install")
+      end
+      cli_server set_pumactl_args, env: env
+    end
+
+    connection = connect("/worker_count")
+    body = read_body(connection, 1)
+    assert_equal(get_stats.fetch("workers").to_s, body)
+  end
+end
diff --git a/test/test_worker_gem_independence.rb b/test/test_worker_gem_independence.rb
@@ -138,13 +138,4 @@ def start_phased_restart
 
     true while @server.gets !~ /booted in [.0-9]+s, phase: 1/
   end
-
-  def with_unbundled_env
-    bundler_ver = Gem::Version.new(Bundler::VERSION)
-    if bundler_ver < Gem::Version.new('2.1.0')
-      Bundler.with_clean_env { yield }
-    else
-      Bundler.with_unbundled_env { yield }
-    end
-  end
 end
diff --git a/test/web_concurrency_test/Gemfile b/test/web_concurrency_test/Gemfile
@@ -0,0 +1,5 @@
+source "https://rubygems.org"
+
+gem 'puma', path: '../..'
+
+gem 'concurrent-ruby', '~> 1.0'
diff --git a/test/web_concurrency_test/config.ru b/test/web_concurrency_test/config.ru
@@ -0,0 +1,5 @@
+map "/worker_count" do
+  run ->(env) {
+    [200, {}, [Concurrent.available_processor_count.to_i.to_s]]
+  }
+end