From 1038404b5b155285a561c150bb4d8494e8071d06 Mon Sep 17 00:00:00 2001
From: MSP-Greg <MSP-Greg@users.noreply.github.com>
Date: Tue, 4 May 2021 21:14:33 -0500
Subject: [PATCH] systemd - fix event firing (#2591)

* systemd - fix event firing

* Update test_integration_systemd.rb

Co-authored-by: Nate Berkopec <nate.berkopec@gmail.com>
---
 lib/puma/cluster.rb              | 12 ++++++++-
 test/test_integration_systemd.rb | 45 +++++++++++++++++++++++---------
 2 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/lib/puma/cluster.rb b/lib/puma/cluster.rb
index cf393b40bc..907611af22 100644
--- a/lib/puma/cluster.rb
+++ b/lib/puma/cluster.rb
@@ -43,6 +43,7 @@ def stop_workers
     end
 
     def start_phased_restart
+      @events.fire_on_restart!
       @phase += 1
       log "- Starting phased worker restart, phase: #{@phase}"
 
@@ -317,7 +318,7 @@ def setup_signals
 
           stop_workers
           stop
-
+          @events.fire_on_stopped!
           raise(SignalException, "SIGTERM") if @options[:raise_exception_on_sigterm]
           exit 0 # Clean exit, workers were stopped
         end
@@ -411,12 +412,16 @@ def run
 
       begin
         booted = false
+        in_phased_restart = false
+        workers_not_booted = @options[:workers]
 
         while @status == :run
           begin
             if @phased_restart
               start_phased_restart
               @phased_restart = false
+              in_phased_restart = true
+              workers_not_booted = @options[:workers]
             end
 
             check_workers
@@ -444,6 +449,7 @@ def run
                   w.boot!
                   log "- Worker #{w.index} (PID: #{pid}) booted in #{w.uptime.round(2)}s, phase: #{w.phase}"
                   @next_check = Time.now
+                  workers_not_booted -= 1
                 when "e"
                   # external term, see worker method, Signal.trap "SIGTERM"
                   w.instance_variable_set :@term, true
@@ -461,6 +467,10 @@ def run
                 log "! Out-of-sync worker list, no #{pid} worker"
               end
             end
+            if in_phased_restart && workers_not_booted.zero?
+              @events.fire_on_booted!
+              in_phased_restart = false
+            end
 
           rescue Interrupt
             @status = :stop
diff --git a/test/test_integration_systemd.rb b/test/test_integration_systemd.rb
index c3257bc292..151a2e077d 100644
--- a/test/test_integration_systemd.rb
+++ b/test/test_integration_systemd.rb
@@ -31,21 +31,18 @@ def teardown
     ENV["WATCHDOG_USEC"] = nil
   end
 
-  def socket_message
-    @socket.recvfrom(15)[0]
+  def test_systemd_notify_usr1_phased_restart_cluster
+    skip_unless :fork
+    assert_restarts_with_systemd :USR1
   end
 
-  def test_systemd_integration
-    cli_server "test/rackup/hello.ru"
-    assert_equal(socket_message, "READY=1")
-
-    connection = connect
-    restart_server connection
-    assert_equal(socket_message, "RELOADING=1")
-    assert_equal(socket_message, "READY=1")
+  def test_systemd_notify_usr2_hot_restart_cluster
+    skip_unless :fork
+    assert_restarts_with_systemd :USR2
+  end
 
-    stop_server
-    assert_equal(socket_message, "STOPPING=1")
+  def test_systemd_notify_usr2_hot_restart_single
+    assert_restarts_with_systemd :USR2, workers: 0
   end
 
   def test_systemd_watchdog
@@ -59,4 +56,28 @@ def test_systemd_watchdog
     stop_server
     assert_match(socket_message, "STOPPING=1")
   end
+
+  private
+
+  def assert_restarts_with_systemd(signal, workers: 2)
+    cli_server "-w#{workers} test/rackup/hello.ru"
+    assert_equal socket_message, 'READY=1'
+
+    Process.kill signal, @pid
+    connect.write "GET / HTTP/1.1\r\n\r\n"
+    assert_equal socket_message, 'RELOADING=1'
+    assert_equal socket_message, 'READY=1'
+
+    Process.kill signal, @pid
+    connect.write "GET / HTTP/1.1\r\n\r\n"
+    assert_equal socket_message, 'RELOADING=1'
+    assert_equal socket_message, 'READY=1'
+
+    stop_server
+    assert_equal socket_message, 'STOPPING=1'
+  end
+
+  def socket_message
+    @socket.recvfrom(15)[0]
+  end
 end