Improve support for graceful shutdown

- Rename `on_before_phased_restart` to `on_before_graceful_shutdown`
- Support `SIGTERM` on Puma
- Ensure that we allow to run only Puma Cluster mode
- Ensure that we run graceful shutdown only when doing
  a shutdown from the loop
- Support `SIGQUIT` on Unicorn

config/initializers/7_prometheus_metrics.rb

@@ -70,8 +70,8 @@ if defined?(::Unicorn) || defined?(::Puma)
     Gitlab::Metrics::Exporter::WebExporter.instance.start
   end
 
-  Gitlab::Cluster::LifecycleEvents.on_before_phased_restart do
-    # We need to ensure that before we re-exec server
+  Gitlab::Cluster::LifecycleEvents.on_before_graceful_shutdown do
+    # We need to ensure that before we re-exec or shutdown server
     # we do stop the exporter
     Gitlab::Metrics::Exporter::WebExporter.instance.stop
   end

config/initializers/validate_puma.rb

+# frozen_string_literal: true
+
+if defined?(::Puma) && ::Puma.cli_config.options[:workers].to_i.zero?
+  raise 'Puma is only supported in Cluster-mode: workers > 0'
+end

lib/gitlab/cluster/lifecycle_events.rb

@@ -33,7 +33,7 @@ module Gitlab
     #
     #     Sidekiq/Puma Single: This is called immediately.
     #
-    # - on_before_phased_restart:
+    # - on_before_graceful_shutdown:
     #
     #     Unicorn/Puma Cluster: This will be called before a graceful
     #       shutdown of workers starts happening.
@@ -75,9 +75,9 @@ module Gitlab
         end
 
         # Read the config/initializers/cluster_events_before_phased_restart.rb
-        def on_before_phased_restart(&block)
+        def on_before_graceful_shutdown(&block)
           # Defer block execution
-          (@master_phased_restart ||= []) << block
+          (@master_graceful_shutdown ||= []) << block
         end
 
         def on_before_master_restart(&block)
@@ -108,8 +108,8 @@ module Gitlab
           end
         end
 
-        def do_before_phased_restart
-          @master_phased_restart&.each do |block|
+        def do_before_graceful_shutdown
+          @master_graceful_shutdown&.each do |block|
             block.call
           end
         end

lib/gitlab/cluster/mixins/puma_cluster.rb

@@ -8,8 +8,12 @@ module Gitlab
           raise 'missing method Puma::Cluster#stop_workers' unless base.method_defined?(:stop_workers)
         end
 
+        # This looks at internal status of `Puma::Cluster`
+        # https://github.com/puma/puma/blob/v3.12.1/lib/puma/cluster.rb#L333
         def stop_workers
-          Gitlab::Cluster::LifecycleEvents.do_before_phased_restart
+          if @status == :stop # rubocop:disable Gitlab/ModuleWithInstanceVariables
+            Gitlab::Cluster::LifecycleEvents.do_before_graceful_shutdown
+          end
 
           super
         end

lib/gitlab/cluster/mixins/unicorn_http_server.rb

@@ -5,11 +5,26 @@ module Gitlab
     module Mixins
       module UnicornHttpServer
         def self.prepended(base)
-          raise 'missing method Unicorn::HttpServer#reexec' unless base.method_defined?(:reexec)
+          unless base.method_defined?(:reexec) && base.method_defined?(:stop)
+            raise 'missing method Unicorn::HttpServer#reexec or Unicorn::HttpServer#stop'
+          end
         end
 
         def reexec
-          Gitlab::Cluster::LifecycleEvents.do_before_phased_restart
+          Gitlab::Cluster::LifecycleEvents.do_before_graceful_shutdown
+
+          super
+        end
+
+        # The stop on non-graceful shutdown is executed twice:
+        # `#stop(false)` and `#stop`.
+        #
+        # The first stop will wipe-out all workers, so we need to check
+        # the flag and a list of workers
+        def stop(graceful = true)
+          if graceful && @workers.any? # rubocop:disable Gitlab/ModuleWithInstanceVariables
+            Gitlab::Cluster::LifecycleEvents.do_before_graceful_shutdown
+          end
 
           super
         end

spec/lib/gitlab/cluster/mixins/puma_cluster_spec.rb

@@ -8,15 +8,28 @@ describe Gitlab::Cluster::Mixins::PumaCluster do
   PUMA_STARTUP_TIMEOUT = 30
 
   context 'when running Puma in Cluster-mode' do
-    %i[USR1 USR2 INT HUP].each do |signal|
-      it "for #{signal} does execute phased restart block" do
+    using RSpec::Parameterized::TableSyntax
+
+    where(:signal, :exitstatus, :termsig) do
+      # executes phased restart block
+      :USR1 | 140 | nil
+      :USR2 | 140 | nil
+      :INT | 140 | nil
+      :HUP | 140 | nil
+
+      # does not execute phased restart block
+      :TERM | nil | 15
+    end
+
+    with_them do
+      it 'properly handles process lifecycle' do
         with_puma(workers: 1) do |pid|
           Process.kill(signal, pid)
 
           child_pid, child_status = Process.wait2(pid)
           expect(child_pid).to eq(pid)
-          expect(child_status).to be_exited
-          expect(child_status.exitstatus).to eq(140)
+          expect(child_status.exitstatus).to eq(exitstatus)
+          expect(child_status.termsig).to eq(termsig)
         end
       end
     end
@@ -62,9 +75,13 @@ describe Gitlab::Cluster::Mixins::PumaCluster do
 
         Puma::Cluster.prepend(#{described_class})
 
-        Gitlab::Cluster::LifecycleEvents.on_before_phased_restart do
+        mutex = Mutex.new
+
+        Gitlab::Cluster::LifecycleEvents.on_before_graceful_shutdown do
+          mutex.synchronize do
             exit(140)
           end
+        end
 
         # redirect stderr to stdout
         $stderr.reopen($stdout)

spec/lib/gitlab/cluster/mixins/unicorn_http_server_spec.rb

@@ -5,31 +5,30 @@ require 'spec_helper'
 # For easier debugging set `UNICORN_DEBUG=1`
 
 describe Gitlab::Cluster::Mixins::UnicornHttpServer do
-  UNICORN_STARTUP_TIMEOUT = 10
+  UNICORN_STARTUP_TIMEOUT = 30
 
   context 'when running Unicorn' do
-    %i[USR2].each do |signal|
-      it "for #{signal} does execute phased restart block" do
-        with_unicorn(workers: 1) do |pid|
-          Process.kill(signal, pid)
+    using RSpec::Parameterized::TableSyntax
 
-          child_pid, child_status = Process.wait2(pid)
-          expect(child_pid).to eq(pid)
-          expect(child_status).to be_exited
-          expect(child_status.exitstatus).to eq(140)
-        end
-      end
+    where(:signal, :exitstatus, :termsig) do
+      # executes phased restart block
+      :USR2 | 140 | nil
+      :QUIT | 140 | nil
+
+      # does not execute phased restart block
+      :INT | 0 | nil
+      :TERM | 0 | nil
     end
 
-    %i[QUIT TERM INT].each do |signal|
-      it "for #{signal} does not execute phased restart block" do
+    with_them do
+      it 'properly handles process lifecycle' do
         with_unicorn(workers: 1) do |pid|
           Process.kill(signal, pid)
 
           child_pid, child_status = Process.wait2(pid)
           expect(child_pid).to eq(pid)
-          expect(child_status).to be_exited
-          expect(child_status.exitstatus).to eq(0)
+          expect(child_status.exitstatus).to eq(exitstatus)
+          expect(child_status.termsig).to eq(termsig)
         end
       end
     end
@@ -74,9 +73,13 @@ describe Gitlab::Cluster::Mixins::UnicornHttpServer do
 
         Unicorn::HttpServer.prepend(#{described_class})
 
-        Gitlab::Cluster::LifecycleEvents.on_before_phased_restart do
+        mutex = Mutex.new
+
+        Gitlab::Cluster::LifecycleEvents.on_before_graceful_shutdown do
+          mutex.synchronize do
             exit(140)
           end
+        end
 
         # redirect stderr to stdout
         $stderr.reopen($stdout)