Add latest changes from gitlab-org/gitlab@master

1 files changed, 50 insertions, 11 deletions

diff --git a/lib/gitlab/process_supervisor.rb b/lib/gitlab/process_supervisor.rb
index f0d2bbc33bd..18fd24aa582 100644
--- a/lib/gitlab/process_supervisor.rb
+++ b/lib/gitlab/process_supervisor.rb
@@ -9,7 +9,7 @@ module Gitlab
   # The supervisor will also trap termination signals if provided and
   # propagate those to the supervised processes. Any supervised processes
   # that do not terminate within a specified grace period will be killed.
-  class ProcessSupervisor
+  class ProcessSupervisor < Gitlab::Daemon
     DEFAULT_HEALTH_CHECK_INTERVAL_SECONDS = 5
     DEFAULT_TERMINATE_INTERVAL_SECONDS = 1
     DEFAULT_TERMINATE_TIMEOUT_SECONDS = 10
@@ -21,13 +21,18 @@ module Gitlab
       check_terminate_interval_seconds: DEFAULT_TERMINATE_INTERVAL_SECONDS,
       terminate_timeout_seconds: DEFAULT_TERMINATE_TIMEOUT_SECONDS,
       term_signals: %i(INT TERM),
-      forwarded_signals: [])
+      forwarded_signals: [],
+      **options)
+      super(**options)
 
       @term_signals = term_signals
       @forwarded_signals = forwarded_signals
       @health_check_interval_seconds = health_check_interval_seconds
       @check_terminate_interval_seconds = check_terminate_interval_seconds
       @terminate_timeout_seconds = terminate_timeout_seconds
+
+      @pids = []
+      @alive = false
     end
 
     # Starts a supervision loop for the given process ID(s).
@@ -39,32 +44,66 @@ module Gitlab
     # start observing those processes instead. Otherwise it will shut down.
     def supervise(pid_or_pids, &on_process_death)
       @pids = Array(pid_or_pids)
+      @on_process_death = on_process_death
 
       trap_signals!
 
+      start
+    end
+
+    # Shuts down the supervisor and all supervised processes with the given signal.
+    def shutdown(signal = :TERM)
+      return unless @alive
+
+      stop_processes(signal)
+      stop
+    end
+
+    def supervised_pids
+      @pids
+    end
+
+    private
+
+    def start_working
       @alive = true
+    end
+
+    def stop_working
+      @alive = false
+    end
+
+    def run_thread
       while @alive
         sleep(@health_check_interval_seconds)
 
-        check_process_health(&on_process_death)
+        check_process_health
       end
     end
 
-    private
-
-    def check_process_health(&on_process_death)
+    def check_process_health
       unless all_alive?
-        dead_pids = @pids - live_pids
-        @pids = Array(yield(dead_pids))
+        existing_pids = live_pids # Capture this value for the duration of the block.
+        dead_pids = @pids - existing_pids
+        new_pids = Array(@on_process_death.call(dead_pids))
+        @pids = existing_pids + new_pids
         @alive = @pids.any?
       end
     end
 
+    def stop_processes(signal)
+      # Set this prior to shutting down so that shutdown hooks which read `alive`
+      # know the supervisor is about to shut down.
+      @alive = false
+
+      # Shut down supervised processes.
+      signal_all(signal)
+      wait_for_termination
+    end
+
     def trap_signals!
       ProcessManagement.trap_signals(@term_signals) do |signal|
-        @alive = false
-        signal_all(signal)
-        wait_for_termination
+        stop_processes(signal)
       end
 
       ProcessManagement.trap_signals(@forwarded_signals) do |signal|