Add `on_before_phased_restart` lifecycle event

Unicorn/Puma Cluster: This will be called before a graceful
shutdown of workers starts hapepning.
This is called on `master` process.

config/initializers/7_prometheus_metrics.rb

@@ -70,9 +70,18 @@ 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
+    # we do stop the exporter
+    Gitlab::Metrics::Exporter::WebExporter.instance.stop
+  end
+
   Gitlab::Cluster::LifecycleEvents.on_before_master_restart do
     # We need to ensure that before we re-exec server
     # we do stop the exporter
+    #
+    # We do it again, for being extra safe,
+    # but it should not be needed
     Gitlab::Metrics::Exporter::WebExporter.instance.stop
   end
 
@@ -81,7 +90,6 @@ if defined?(::Unicorn) || defined?(::Puma)
     # but this does not happen for Ruby fork
     #
     # This does stop server, as it is running on master.
-    # However, ensures that we close the TCPSocket.
     Gitlab::Metrics::Exporter::WebExporter.instance.stop
   end
 end

config/initializers/cluster_events_before_phased_restart.rb

+# Technical debt, this should be ideally upstreamed.
+#
+# However, there's currently no way to hook before doing
+# graceful shutdown today.
+#
+# Follow-up the issue: https://gitlab.com/gitlab-org/gitlab/issues/34107
+
+if defined?(::Puma)
+  Puma::Cluster.prepend(::Gitlab::Cluster::Mixins::PumaCluster)
+end
+
+if defined?(::Unicorn::HttpServer)
+  Unicorn::HttpServer.prepend(::Gitlab::Cluster::Mixins::UnicornHttpServer)
+end

lib/gitlab/cluster/lifecycle_events.rb

@@ -8,14 +8,50 @@ module Gitlab
     # watchdog threads. This lets us abstract away the Unix process
     # lifecycles of Unicorn, Sidekiq, Puma, Puma Cluster, etc.
     #
-    # We have three lifecycle events.
+    # We have the following lifecycle events.
     #
-    # - before_fork (only in forking processes)
-    #     In forking processes (Unicorn and Puma in multiprocess mode) this
-    #     will be called exactly once, on startup, before the workers are
-    #     forked. This will be called in the parent process.
-    # - worker_start
-    # - before_master_restart (only in forking processes)
+    # - on_master_start:
+    #
+    #     Unicorn/Puma Cluster: This will be called exactly once,
+    #       on startup, before the workers are forked. This is
+    #       called in the PARENT/MASTER process.
+    #
+    #     Sidekiq/Puma Single: This is called immediately.
+    #
+    # - on_before_fork:
+    #
+    #     Unicorn/Puma Cluster: This will be called exactly once,
+    #       on startup, before the workers are forked. This is
+    #       called in the PARENT/MASTER process.
+    #
+    #     Sidekiq/Puma Single: This is not called.
+    #
+    # - on_worker_start:
+    #
+    #     Unicorn/Puma Cluster: This is called in the worker process
+    #       exactly once before processing requests.
+    #
+    #     Sidekiq/Puma Single: This is called immediately.
+    #
+    # - on_before_phased_restart:
+    #
+    #     Unicorn/Puma Cluster: This will be called before a graceful
+    #       shutdown of workers starts happening.
+    #       This is called on `master` process.
+    #
+    #     Sidekiq/Puma Single: This is not called.
+    #
+    # - on_before_master_restart:
+    #
+    #     Unicorn: This will be called before a new master is spun up.
+    #       This is called on forked master before `execve` to become
+    #       a new masterfor Unicorn. This means that this does not really
+    #       affect old master process.
+    #
+    #     Puma Cluster: This will be called before a new master is spun up.
+    #       This is called on `master` process.
+    #
+    #     Sidekiq/Puma Single: This is not called.
     #
     # Blocks will be executed in the order in which they are registered.
     #
@@ -34,15 +70,17 @@ module Gitlab
         end
 
         def on_before_fork(&block)
-          return unless in_clustered_environment?
-
           # Defer block execution
           (@before_fork_hooks ||= []) << block
         end
 
-        def on_before_master_restart(&block)
-          return unless in_clustered_environment?
+        # Read the config/initializers/cluster_events_before_phased_restart.rb
+        def on_before_phased_restart(&block)
+          # Defer block execution
+          (@master_phased_restart ||= []) << block
+        end
 
+        def on_before_master_restart(&block)
           # Defer block execution
           (@master_restart_hooks ||= []) << block
         end
@@ -70,8 +108,14 @@ module Gitlab
           end
         end
 
+        def do_before_phased_restart
+          @master_phased_restart&.each do |block|
+            block.call
+          end
+        end
+
         def do_before_master_restart
-          @master_restart_hooks && @master_restart_hooks.each do |block|
+          @master_restart_hooks&.each do |block|
             block.call
           end
         end

lib/gitlab/cluster/mixins/puma_cluster.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Cluster
+    module Mixins
+      module PumaCluster
+        def self.prepended(base)
+          raise 'missing method Puma::Cluster#stop_workers' unless base.method_defined?(:stop_workers)
+        end
+
+        def stop_workers
+          Gitlab::Cluster::LifecycleEvents.do_before_phased_restart
+
+          super
+        end
+      end
+    end
+  end
+end

lib/gitlab/cluster/mixins/unicorn_http_server.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Cluster
+    module Mixins
+      module UnicornHttpServer
+        def self.prepended(base)
+          raise 'missing method Unicorn::HttpServer#reexec' unless base.method_defined?(:reexec)
+        end
+
+        def reexec
+          Gitlab::Cluster::LifecycleEvents.do_before_phased_restart
+
+          super
+        end
+      end
+    end
+  end
+end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# For easier debugging set `PUMA_DEBUG=1`
+
+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
+        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)
+        end
+      end
+    end
+  end
+
+  private
+
+  def with_puma(workers:, timeout: PUMA_STARTUP_TIMEOUT)
+    with_puma_config(workers: workers) do |puma_rb|
+      cmdline = [
+        "bundle", "exec", "puma",
+        "-C", puma_rb,
+        "-I", Rails.root.to_s
+      ]
+
+      IO.popen(cmdline) do |process|
+        # wait for process to start:
+        # [2123] * Listening on tcp://127.0.0.1:0
+        wait_for_output(process, /Listening on/, timeout: timeout)
+        consume_output(process)
+
+        yield(process.pid)
+      ensure
+        Process.kill(:KILL, process.pid) unless process.eof?
+      end
+    end
+  end
+
+  def with_puma_config(workers:)
+    Dir.mktmpdir do |dir|
+      File.write "#{dir}/puma.rb", <<-EOF
+        require './lib/gitlab/cluster/lifecycle_events'
+        require './lib/gitlab/cluster/mixins/puma_cluster'
+
+        workers #{workers}
+        bind "tcp://127.0.0.1:0"
+        preload_app!
+
+        app -> (env) { [404, {}, ['']] }
+
+        Puma::Cluster.prepend(#{described_class})
+
+        Gitlab::Cluster::LifecycleEvents.on_before_phased_restart do
+          exit(140)
+        end
+
+        # redirect stderr to stdout
+        $stderr.reopen($stdout)
+      EOF
+
+      yield("#{dir}/puma.rb")
+    end
+  end
+
+  def wait_for_output(process, output, timeout:)
+    Timeout.timeout(timeout) do
+      loop do
+        line = process.readline
+        puts "PUMA_DEBUG: #{line}" if ENV['PUMA_DEBUG']
+        break if line =~ output
+      end
+    end
+  end
+
+  def consume_output(process)
+    Thread.new do
+      loop do
+        line = process.readline
+        puts "PUMA_DEBUG: #{line}" if ENV['PUMA_DEBUG']
+      end
+    rescue
+    end
+  end
+end

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

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+# For easier debugging set `UNICORN_DEBUG=1`
+
+describe Gitlab::Cluster::Mixins::UnicornHttpServer do
+  UNICORN_STARTUP_TIMEOUT = 10
+
+  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)
+
+          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
+    end
+
+    %i[QUIT TERM INT].each do |signal|
+      it "for #{signal} does not execute phased restart block" 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)
+        end
+      end
+    end
+  end
+
+  private
+
+  def with_unicorn(workers:, timeout: UNICORN_STARTUP_TIMEOUT)
+    with_unicorn_configs(workers: workers) do |unicorn_rb, config_ru|
+      cmdline = [
+        "bundle", "exec", "unicorn",
+        "-I", Rails.root.to_s,
+        "-c", unicorn_rb,
+        config_ru
+      ]
+
+      IO.popen(cmdline) do |process|
+        # wait for process to start:
+        # I, [2019-10-15T13:21:27.565225 #3089]  INFO -- : master process ready
+        wait_for_output(process, /master process ready/, timeout: timeout)
+        consume_output(process)
+
+        yield(process.pid)
+      ensure
+        Process.kill(:KILL, process.pid) unless process.eof?
+      end
+    end
+  end
+
+  def with_unicorn_configs(workers:)
+    Dir.mktmpdir do |dir|
+      File.write "#{dir}/unicorn.rb", <<-EOF
+        require './lib/gitlab/cluster/lifecycle_events'
+        require './lib/gitlab/cluster/mixins/unicorn_http_server'
+
+        worker_processes #{workers}
+        listen "127.0.0.1:0"
+        preload_app true
+
+        Unicorn::HttpServer.prepend(#{described_class})
+
+        Gitlab::Cluster::LifecycleEvents.on_before_phased_restart do
+          exit(140)
+        end
+
+        # redirect stderr to stdout
+        $stderr.reopen($stdout)
+      EOF
+
+      File.write "#{dir}/config.ru", <<-EOF
+        run -> (env) { [404, {}, ['']] }
+      EOF
+
+      yield("#{dir}/unicorn.rb", "#{dir}/config.ru")
+    end
+  end
+
+  def wait_for_output(process, output, timeout:)
+    Timeout.timeout(timeout) do
+      loop do
+        line = process.readline
+        puts "UNICORN_DEBUG: #{line}" if ENV['UNICORN_DEBUG']
+        break if line =~ output
+      end
+    end
+  end
+
+  def consume_output(process)
+    Thread.new do
+      loop do
+        line = process.readline
+        puts "UNICORN_DEBUG: #{line}" if ENV['UNICORN_DEBUG']
+      end
+    rescue
+    end
+  end
+end