From 06e88c1d85d0ff91898c941f39cd1c62750ef2df Mon Sep 17 00:00:00 2001
From: Olivier Bellone <ob@stripe.com>
Date: Sat, 1 Jan 2022 12:55:57 -0800
Subject: [PATCH] Allow culling of oldest workers, previously was only youngest
 (#2773)

---
 lib/puma/cluster.rb              | 17 +++++++----
 lib/puma/configuration.rb        |  1 +
 lib/puma/dsl.rb                  | 24 +++++++++++++++
 test/test_integration_cluster.rb | 50 ++++++++++++++++++++++++++++++++
 4 files changed, 87 insertions(+), 5 deletions(-)

diff --git a/lib/puma/cluster.rb b/lib/puma/cluster.rb
index 7d6d12189e..d90c81d569 100644
--- a/lib/puma/cluster.rb
+++ b/lib/puma/cluster.rb
@@ -111,7 +111,14 @@ def cull_workers
 
       debug "Culling #{diff.inspect} workers"
 
-      workers_to_cull = @workers[-diff,diff]
+      workers_to_cull =
+        case @options[:worker_culling_strategy]
+        when :youngest
+          @workers.sort_by(&:started_at)[-diff,diff]
+        when :oldest
+          @workers.sort_by(&:started_at)[0,diff]
+        end
+
       debug "Workers to cull: #{workers_to_cull.inspect}"
 
       workers_to_cull.each do |worker|
@@ -122,10 +129,10 @@ def cull_workers
 
     # @!attribute [r] next_worker_index
     def next_worker_index
-      all_positions =  0...@options[:workers]
-      occupied_positions = @workers.map { |w| w.index }
-      available_positions = all_positions.to_a - occupied_positions
-      available_positions.first
+      occupied_positions = @workers.map(&:index)
+      idx = 0
+      idx += 1 until !occupied_positions.include?(idx)
+      idx
     end
 
     def all_workers_booted?
diff --git a/lib/puma/configuration.rb b/lib/puma/configuration.rb
index 974709b097..9871ff0212 100644
--- a/lib/puma/configuration.rb
+++ b/lib/puma/configuration.rb
@@ -200,6 +200,7 @@ def puma_default_options
         :worker_timeout => DefaultWorkerTimeout,
         :worker_boot_timeout => DefaultWorkerTimeout,
         :worker_shutdown_timeout => DefaultWorkerShutdownTimeout,
+        :worker_culling_strategy => :youngest,
         :remote_address => :socket,
         :tag => method(:infer_tag),
         :environment => -> { ENV['APP_ENV'] || ENV['RACK_ENV'] || ENV['RAILS_ENV'] || 'development' },
diff --git a/lib/puma/dsl.rb b/lib/puma/dsl.rb
index 308a9687fe..a221587b1d 100644
--- a/lib/puma/dsl.rb
+++ b/lib/puma/dsl.rb
@@ -795,6 +795,30 @@ def worker_shutdown_timeout(timeout)
       @options[:worker_shutdown_timeout] = Integer(timeout)
     end
 
+    # Set the strategy for worker culling.
+    #
+    # There are two possible values:
+    #
+    # 1. **:youngest** - the youngest workers (i.e. the workers that were
+    #    the most recently started) will be culled.
+    # 2. **:oldest** - the oldest workers (i.e. the workers that were started
+    #    the longest time ago) will be culled.
+    #
+    # @note Cluster mode only.
+    # @example
+    #   worker_culling_strategy :oldest
+    # @see Puma::Cluster#cull_workers
+    #
+    def worker_culling_strategy(strategy)
+      stategy = strategy.to_sym
+
+      if ![:youngest, :oldest].include?(strategy)
+        raise "Invalid value for worker_culling_strategy - #{stategy}"
+      end
+
+      @options[:worker_culling_strategy] = strategy
+    end
+
     # When set to true (the default), workers accept all requests
     # and queue them before passing them to the handlers.
     # When set to false, each worker process accepts exactly as
diff --git a/test/test_integration_cluster.rb b/test/test_integration_cluster.rb
index 0db66612b8..aaccd38507 100644
--- a/test/test_integration_cluster.rb
+++ b/test/test_integration_cluster.rb
@@ -351,6 +351,56 @@ def test_warning_message_not_outputted_when_single_worker_silenced
     refute_match(/WARNING: Detected running cluster mode with 1 worker/, output.join)
   end
 
+  def test_signal_ttin
+    cli_server "-w 2 test/rackup/hello.ru"
+    get_worker_pids # to consume server logs
+
+    Process.kill :TTIN, @pid
+
+    line = @server.gets
+    assert_match(/Worker 2 \(PID: \d+\) booted in/, line)
+  end
+
+  def test_signal_ttou
+    cli_server "-w 2 test/rackup/hello.ru"
+    get_worker_pids # to consume server logs
+
+    Process.kill :TTOU, @pid
+
+    line = @server.gets
+    assert_match(/Worker 1 \(PID: \d+\) terminating/, line)
+  end
+
+  def test_culling_strategy_youngest
+    cli_server "-w 2 test/rackup/hello.ru", config: "worker_culling_strategy :youngest"
+    get_worker_pids # to consume server logs
+
+    Process.kill :TTIN, @pid
+
+    line = @server.gets
+    assert_match(/Worker 2 \(PID: \d+\) booted in/, line)
+
+    Process.kill :TTOU, @pid
+
+    line = @server.gets
+    assert_match(/Worker 2 \(PID: \d+\) terminating/, line)
+  end
+
+  def test_culling_strategy_oldest
+    cli_server "-w 2 test/rackup/hello.ru", config: "worker_culling_strategy :oldest"
+    get_worker_pids # to consume server logs
+
+    Process.kill :TTIN, @pid
+
+    line = @server.gets
+    assert_match(/Worker 2 \(PID: \d+\) booted in/, line)
+
+    Process.kill :TTOU, @pid
+
+    line = @server.gets
+    assert_match(/Worker 0 \(PID: \d+\) terminating/, line)
+  end
+
   private
 
   def worker_timeout(timeout, iterations, details, config)