Add lifecycle hook for worker stop event

This will be used soon in order to produce
heap dumps when a worker was terminated due
to high memory use.

config/initializers/sidekiq.rb

@@ -83,6 +83,10 @@ def enable_semi_reliable_fetch_mode?
     end
   end
 
+  config.on(:shutdown) do
+    Gitlab::Cluster::LifecycleEvents.do_worker_stop
+  end
+
   if enable_reliable_fetch?
     config.options[:semi_reliable_fetch] = enable_semi_reliable_fetch_mode?
     Sidekiq::ReliableFetch.setup_reliable_fetch!(config)

config/puma.example.development.rb

@@ -67,6 +67,11 @@
   Gitlab::Cluster::LifecycleEvents.do_worker_start
 end
 
+on_worker_shutdown do
+  # Signal application hooks that a worker is shutting down
+  Gitlab::Cluster::LifecycleEvents.do_worker_stop
+end
+
 # Preload the application before starting the workers; this conflicts with
 # phased restart feature. (off by default)
 

config/puma.rb.example

@@ -57,6 +57,11 @@ on_worker_boot do
   Gitlab::Cluster::LifecycleEvents.do_worker_start
 end
 
+on_worker_shutdown do
+  # Signal application hooks that a worker is shutting down
+  Gitlab::Cluster::LifecycleEvents.do_worker_stop
+end
+
 # Preload the application before starting the workers; this conflicts with
 # phased restart feature. (off by default)
 preload_app!

lib/gitlab/cluster/lifecycle_events.rb

@@ -63,6 +63,15 @@ module Cluster
     #
     #     Sidekiq/Puma Single: This is called immediately.
     #
+    # - on_worker_stop (on worker process):
+    #
+    #     Puma Cluster: Called in the worker process
+    #       exactly once after it stops processing requests
+    #       but before it shuts down.
+    #
+    #     Sidekiq: Called after the scheduler shuts down but
+    #       before the worker finishes ongoing jobs.
+    #
     # Blocks will be executed in the order in which they are registered.
     #
     class LifecycleEvents
@@ -113,6 +122,10 @@ def on_master_start(&block)
           end
         end
 
+        def on_worker_stop(&block)
+          (@worker_stop_hooks ||= []) << block
+        end
+
         #
         # Lifecycle integration methods (called from puma.rb, etc.)
         #
@@ -137,6 +150,10 @@ def do_before_master_restart
           call(:master_restart_hooks, @master_restart_hooks)
         end
 
+        def do_worker_stop
+          call(:worker_stop_hooks, @worker_stop_hooks)
+        end
+
         # DEPRECATED
         alias_method :do_master_restart, :do_before_master_restart
 

spec/lib/gitlab/cluster/lifecycle_events_spec.rb

@@ -3,38 +3,55 @@
 require 'spec_helper'
 
 RSpec.describe Gitlab::Cluster::LifecycleEvents do
+  using RSpec::Parameterized::TableSyntax
+
   # we create a new instance to ensure that we do not touch existing hooks
   let(:replica) { Class.new(described_class) }
 
-  context 'hooks execution' do
-    using RSpec::Parameterized::TableSyntax
+  before do
+    # disable blackout period to speed-up tests
+    stub_config(shutdown: { blackout_seconds: 0 })
+  end
 
-    where(:method, :hook_names) do
-      :do_worker_start             | %i[worker_start_hooks]
-      :do_before_fork              | %i[before_fork_hooks]
-      :do_before_graceful_shutdown | %i[master_blackout_period master_graceful_shutdown]
-      :do_before_master_restart    | %i[master_restart_hooks]
+  context 'outside of clustered environments' do
+    where(:hook, :was_executed_immediately) do
+      :on_worker_start             | true
+      :on_before_fork              | false
+      :on_before_graceful_shutdown | false
+      :on_before_master_restart    | false
+      :on_worker_stop              | false
     end
 
-    before do
-      # disable blackout period to speed-up tests
-      stub_config(shutdown: { blackout_seconds: 0 })
+    with_them do
+      it 'executes the given block immediately' do
+        was_executed = false
+        replica.public_send(hook, &proc { was_executed = true })
+
+        expect(was_executed).to eq(was_executed_immediately)
+      end
     end
+  end
 
-    with_them do
-      subject { replica.public_send(method) }
+  context 'in clustered environments' do
+    before do
+      allow(Gitlab::Runtime).to receive(:puma?).and_return(true)
+      replica.set_puma_options(workers: 2)
+    end
 
-      it 'executes all hooks' do
-        hook_names.each do |hook_name|
-          hook = double
-          replica.instance_variable_set(:"@#{hook_name}", [hook])
+    where(:hook, :execution_helper) do
+      :on_worker_start             | :do_worker_start
+      :on_before_fork              | :do_before_fork
+      :on_before_graceful_shutdown | :do_before_graceful_shutdown
+      :on_before_master_restart    | :do_before_master_restart
+      :on_worker_stop              | :do_worker_stop
+    end
 
-          # ensure that proper hooks are called
-          expect(hook).to receive(:call)
-          expect(replica).to receive(:call).with(hook_name, anything).and_call_original
-        end
+    with_them do
+      it 'requires explicit execution via do_* helper' do
+        was_executed = false
+        replica.public_send(hook, &proc { was_executed = true })
 
-        subject
+        expect { replica.public_send(execution_helper) }.to change { was_executed }.from(false).to(true)
       end
     end
   end