From 6543f0959bd495b8b0326777034d6f13651dd4b3 Mon Sep 17 00:00:00 2001
From: MSP-Greg <MSP-Greg@users.noreply.github.com>
Date: Thu, 2 Jun 2022 13:10:03 -0500
Subject: [PATCH] Puma::ControlCLI - allow refork command to be sent as a
 request (#2868)

* Puma::ControlCLI - allow refork command to be sent as a request

* Puma::ControlCLI - add signal only commands

* Puma::ControlCLI - check whether signal is available
---
 lib/puma/app/status.rb           |  3 +++
 lib/puma/control_cli.rb          | 30 +++++++++++++++----------
 lib/puma/launcher.rb             | 11 +++++++++
 test/test_integration_cluster.rb |  2 +-
 test/test_integration_pumactl.rb | 37 +++++++++++++++++++++++++++++++
 test/test_pumactl.rb             | 38 +++++++++++++++++++++++++++++++-
 6 files changed, 107 insertions(+), 14 deletions(-)

diff --git a/lib/puma/app/status.rb b/lib/puma/app/status.rb
index 5cfb2b0b..7fe903cb 100644
--- a/lib/puma/app/status.rb
+++ b/lib/puma/app/status.rb
@@ -39,6 +39,9 @@ module Puma
           when 'phased-restart'
             @launcher.phased_restart ? 200 : 404
 
+          when 'refork'
+            @launcher.refork ? 200 : 404
+
           when 'reload-worker-directory'
             @launcher.send(:reload_worker_directory) ? 200 : 404
 
diff --git a/lib/puma/control_cli.rb b/lib/puma/control_cli.rb
index 23030962..859e5107 100644
--- a/lib/puma/control_cli.rb
+++ b/lib/puma/control_cli.rb
@@ -17,26 +17,30 @@ module Puma
     CMD_PATH_SIG_MAP = {
       'gc'       => nil,
       'gc-stats' => nil,
-      'halt'     => 'SIGQUIT',
-      'phased-restart' => 'SIGUSR1',
-      'refork'   => 'SIGURG',
+      'halt'              => 'SIGQUIT',
+      'info'              => 'SIGINFO',
+      'phased-restart'    => 'SIGUSR1',
+      'refork'            => 'SIGURG',
       'reload-worker-directory' => nil,
-      'restart'  => 'SIGUSR2',
+      'reopen-log'        => 'SIGHUP',
+      'restart'           => 'SIGUSR2',
       'start'    => nil,
       'stats'    => nil,
       'status'   => '',
-      'stop'     => 'SIGTERM',
-      'thread-backtraces' => nil
+      'stop'              => 'SIGTERM',
+      'thread-backtraces' => nil,
+      'worker-count-down' => 'SIGTTOU',
+      'worker-count-up'   => 'SIGTTIN'
     }.freeze
 
     # @deprecated 6.0.0
     COMMANDS = CMD_PATH_SIG_MAP.keys.freeze
 
     # commands that cannot be used in a request
-    NO_REQ_COMMANDS = %w{refork}.freeze
+    NO_REQ_COMMANDS = %w[info reopen-log worker-count-down worker-count-up].freeze
 
     # @version 5.0.0
-    PRINTABLE_COMMANDS = %w{gc-stats stats thread-backtraces}.freeze
+    PRINTABLE_COMMANDS = %w[gc-stats stats thread-backtraces].freeze
 
     def initialize(argv, stdout=STDOUT, stderr=STDERR)
       @state = nil
@@ -185,8 +189,6 @@ module Puma
 
       if @command == 'status'
         message 'Puma is started'
-      elsif NO_REQ_COMMANDS.include? @command
-        raise "Invalid request command: #{@command}"
       else
         url = "/#{@command}"
 
@@ -242,7 +244,11 @@ module Puma
           @stdout.flush unless @stdout.sync
           return
         elsif sig.start_with? 'SIG'
-          Process.kill sig, @pid
+          if Signal.list.key? sig.sub(/\ASIG/, '')
+            Process.kill sig, @pid
+          else
+            raise "Signal '#{sig}' not available'"
+          end
         elsif @command == 'status'
           begin
             Process.kill 0, @pid
@@ -268,7 +274,7 @@ module Puma
       return start if @command == 'start'
       prepare_configuration
 
-      if Puma.windows? || @control_url
+      if Puma.windows? || @control_url && !NO_REQ_COMMANDS.include?(@command)
         send_request
       else
         send_signal
diff --git a/lib/puma/launcher.rb b/lib/puma/launcher.rb
index 7fc4652d..67e068d3 100644
--- a/lib/puma/launcher.rb
+++ b/lib/puma/launcher.rb
@@ -163,6 +163,17 @@ module Puma
       true
     end
 
+    # Begin a refork if supported
+    def refork
+      if clustered? && @runner.respond_to?(:fork_worker!) && @options[:fork_worker]
+        @runner.fork_worker!
+        true
+      else
+        log "* refork called but not available."
+        false
+      end
+    end
+
     # Run the server. This blocks until the server is stopped
     def run
       previous_env = get_env
diff --git a/test/test_integration_cluster.rb b/test/test_integration_cluster.rb
index fde881ae..7fdd6e98 100644
--- a/test/test_integration_cluster.rb
+++ b/test/test_integration_cluster.rb
@@ -229,7 +229,7 @@ RUBY
   end
 
   # use three workers to keep accepting clients
-  def test_refork
+  def test_fork_worker_on_refork
     refork = Tempfile.new 'refork'
     wrkrs = 3
     cli_server "-w #{wrkrs} test/rackup/hello_with_delay.ru", config: <<RUBY
diff --git a/test/test_integration_pumactl.rb b/test/test_integration_pumactl.rb
index bf90daf5..8840e2d5 100644
--- a/test/test_integration_pumactl.rb
+++ b/test/test_integration_pumactl.rb
@@ -93,6 +93,43 @@ class TestIntegrationPumactl < TestIntegration
     @server = nil
   end
 
+  def test_refork_cluster
+    skip_unless :fork
+    wrkrs = 3
+    cli_server "-q -w #{wrkrs} test/rackup/sleep.ru --control-url unix://#{@control_path} --control-token #{TOKEN} -S #{@state_path}",
+      config: 'fork_worker 50',
+      unix: true
+
+    start = Time.now
+
+    s = UNIXSocket.new @bind_path
+    @ios_to_close << s
+    s << "GET /sleep1 HTTP/1.0\r\n\r\n"
+
+    # Get the PIDs of the phase 0 workers.
+    phase0_worker_pids = get_worker_pids 0, wrkrs
+    assert File.exist? @bind_path
+
+    cli_pumactl "refork", unix: true
+
+    # Get the PIDs of the phase 1 workers.
+    phase1_worker_pids = get_worker_pids 1, wrkrs - 1
+
+    msg = "phase 0 pids #{phase0_worker_pids.inspect}  phase 1 pids #{phase1_worker_pids.inspect}"
+
+    assert_equal wrkrs    , phase0_worker_pids.length, msg
+    assert_equal wrkrs - 1, phase1_worker_pids.length, msg
+    assert_empty phase0_worker_pids & phase1_worker_pids, "#{msg}\nBoth workers should be replaced with new"
+    assert File.exist?(@bind_path), "Bind path must exist after phased refork"
+
+    cli_pumactl "stop", unix: true
+
+    _, status = Process.wait2(@pid)
+    assert_equal 0, status
+    assert_operator Time.now - start, :<, (DARWIN ? 8 : 6)
+    @server = nil
+  end
+
   def test_prune_bundler_with_multiple_workers
     skip_unless :fork
 
diff --git a/test/test_pumactl.rb b/test/test_pumactl.rb
index 156b647c..f1ed5862 100644
--- a/test/test_pumactl.rb
+++ b/test/test_pumactl.rb
@@ -19,7 +19,7 @@ class TestPumaControlCli < TestConfigFileBase
 
   def teardown
     @wait.close
-    @ready.close
+    @ready.close unless @ready.closed?
   end
 
   def with_config_file(path_to_config, port)
@@ -187,6 +187,42 @@ class TestPumaControlCli < TestConfigFileBase
     assert_kind_of Thread, t.join, "server didn't stop"
   end
 
+  # This checks that a 'signal only' command is sent
+  # they are defined by the `Puma::ControlCLI::NO_REQ_COMMANDS` array
+  # test is skipped unless NO_REQ_COMMANDS is defined
+  def test_control_url_with_signal_only_cmd
+    skip_if :windows
+    skip unless defined? Puma::ControlCLI::NO_REQ_COMMANDS
+    host = "127.0.0.1"
+    port = UniquePort.call
+    url = "tcp://#{host}:#{port}/"
+
+    opts = [
+      "--control-url", url,
+      "--control-token", "ctrl",
+      "--config-file", "test/config/app.rb",
+      "--pid", "1234"
+    ]
+    cmd = Puma::ControlCLI::NO_REQ_COMMANDS.first
+    log = ''.dup
+    control_cli = Puma::ControlCLI.new (opts + [cmd]), @ready, @ready
+
+    def control_cli.send_signal
+      message "send_signal #{@command}\n"
+    end
+    def control_cli.send_request
+      message "send_request #{@command}\n"
+    end
+
+    control_cli.run
+    @ready.close
+
+    log = @wait.read
+
+    assert_includes log, "send_signal #{cmd}"
+    refute_includes log, 'send_request'
+  end
+
   def test_control_ssl
     skip_unless :ssl