systemd - fix event firing

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

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
+    restart :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
+    restart :USR2
+  end
 
-    stop_server
-    assert_equal(socket_message, "STOPPING=1")
+  def test_systemd_notify_usr2_hot_restart_single
+      restart :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 restart(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