Introduce and use `refresh_state`

This updates metrics on demand,
and tracks the current state as gauge.

This also improves tests execution time.

config/initializers/7_prometheus_metrics.rb

@@ -41,9 +41,6 @@ Sidekiq.configure_server do |config|
     # after all workers have forked, but I don't know how at this point.
     ::Prometheus::Client.reinitialize_on_pid_change(force: true)
 
-    # temporary solution before fix https://gitlab.com/gitlab-org/gitlab/issues/33125
-    ::Prometheus::Client.reinitialize_on_pid_change(force: true)
-
     Gitlab::Metrics::Exporter::SidekiqExporter.instance.start
   end
 end

lib/gitlab/sidekiq_daemon/memory_killer.rb

@@ -20,6 +20,7 @@ module Gitlab
       # Developer/admin should always set `memory_killer_max_memory_growth_kb` explicitly
       # In case not set, default to 300M. This is for extra-safe.
       DEFAULT_MAX_MEMORY_GROWTH_KB = 300_000
+
       # Phases of memory killer
       PHASE = {
         running: 1,
@@ -40,13 +41,26 @@ module Gitlab
 
       def init_metrics
         {
-          sidekiq_current_rss:    ::Gitlab::Metrics.gauge(:sidekiq_current_rss, 'Current RSS of Sidekiq Worker'),
+          sidekiq_current_rss:                  ::Gitlab::Metrics.gauge(:sidekiq_current_rss, 'Current RSS of Sidekiq Worker'),
           sidekiq_memory_killer_soft_limit_rss: ::Gitlab::Metrics.gauge(:sidekiq_memory_killer_soft_limit_rss, 'Current soft_limit_rss of Sidekiq Worker'),
           sidekiq_memory_killer_hard_limit_rss: ::Gitlab::Metrics.gauge(:sidekiq_memory_killer_hard_limit_rss, 'Current hard_limit_rss of Sidekiq Worker'),
-          sidekiq_memory_killer_phase:  ::Gitlab::Metrics.gauge(:sidekiq_memory_killer_phase, 'Current phase of Sidekiq Worker')
+          sidekiq_memory_killer_phase:          ::Gitlab::Metrics.gauge(:sidekiq_memory_killer_phase, 'Current phase of Sidekiq Worker')
         }
       end
 
+      def refresh_state(phase)
+        @phase = PHASE.fetch(phase)
+        @current_rss = get_rss
+        @soft_limit_rss = get_soft_limit_rss
+        @hard_limit_rss = get_hard_limit_rss
+
+        # track the current state as prometheus gauges
+        @metrics[:sidekiq_memory_killer_phase].set({}, @phase)
+        @metrics[:sidekiq_current_rss].set({}, @current_rss)
+        @metrics[:sidekiq_memory_killer_soft_limit_rss].set({}, @soft_limit_rss)
+        @metrics[:sidekiq_memory_killer_hard_limit_rss].set({}, @hard_limit_rss)
+      end
+
       def run_thread
         Sidekiq.logger.info(
           class: self.class.to_s,
@@ -95,64 +109,55 @@ module Gitlab
         # Tell Sidekiq to stop fetching new jobs
         # We first SIGNAL and then wait given time
         # We also monitor a number of running jobs and allow to restart early
-        update_metrics(PHASE[:stop_fetching_new_jobs], get_rss, get_soft_limit_rss, get_hard_limit_rss)
+        refresh_state(:stop_fetching_new_jobs)
         signal_and_wait(SHUTDOWN_TIMEOUT_SECONDS, 'SIGTSTP', 'stop fetching new jobs')
         return unless enabled?
 
         # Tell sidekiq to restart itself
         # Keep extra safe to wait `Sidekiq.options[:timeout] + 2` seconds before SIGKILL
-        update_metrics(PHASE[:shutting_down], get_rss, get_soft_limit_rss, get_hard_limit_rss)
+        refresh_state(:shutting_down)
         signal_and_wait(Sidekiq.options[:timeout] + 2, 'SIGTERM', 'gracefully shut down')
         return unless enabled?
 
         # Ideally we should never reach this condition
         # Wait for Sidekiq to shutdown gracefully, and kill it if it didn't
         # Kill the whole pgroup, so we can be sure no children are left behind
-        update_metrics(PHASE[:killing_sidekiq], get_rss, get_soft_limit_rss, get_hard_limit_rss)
+        refresh_state(:killing_sidekiq)
         signal_pgroup('SIGKILL', 'die')
       end
 
       def rss_within_range?
-        phase = PHASE[:running]
-        current_rss = nil
-        soft_limit_rss = nil
-        hard_limit_rss = nil
+        refresh_state(:running)
+
         deadline = Gitlab::Metrics::System.monotonic_time + GRACE_BALLOON_SECONDS.seconds
         loop do
           return true unless enabled?
 
-          current_rss = get_rss
-          soft_limit_rss = get_soft_limit_rss
-          hard_limit_rss = get_hard_limit_rss
-
-          update_metrics(phase, current_rss, soft_limit_rss, hard_limit_rss)
-
           # RSS go above hard limit should trigger forcible shutdown right away
-          break if current_rss > hard_limit_rss
+          break if @current_rss > @hard_limit_rss
 
           # RSS go below the soft limit
-          return true if current_rss < soft_limit_rss
-
-          phase = PHASE[:above_soft_limit]
+          return true if @current_rss < @soft_limit_rss
 
           # RSS did not go below the soft limit within deadline, restart
           break if Gitlab::Metrics::System.monotonic_time > deadline
 
           sleep(CHECK_INTERVAL_SECONDS)
+
+          refresh_state(:above_soft_limit)
         end
 
-        log_rss_out_of_range(current_rss, hard_limit_rss, soft_limit_rss)
+        # There are two chances to break from loop:
+        #   - above hard limit, or
+        #   - above soft limit after deadline
+        # When `above hard limit`, it immediately go to `stop_fetching_new_jobs`
+        # So ignore `above hard limit` and always set `above_soft_limit` here
+        refresh_state(:above_soft_limit)
+        log_rss_out_of_range(@current_rss, @hard_limit_rss, @soft_limit_rss)
 
         false
       end
 
-      def update_metrics(phase, current_rss, soft_limit_rss, hard_limit_rss)
-        @metrics[:sidekiq_memory_killer_phase].set({}, phase)
-        @metrics[:sidekiq_current_rss].set({}, current_rss)
-        @metrics[:sidekiq_memory_killer_soft_limit_rss].set({}, soft_limit_rss)
-        @metrics[:sidekiq_memory_killer_hard_limit_rss].set({}, hard_limit_rss)
-      end
-
       def log_rss_out_of_range(current_rss, hard_limit_rss, soft_limit_rss)
         Sidekiq.logger.warn(
           class: self.class.to_s,

spec/lib/gitlab/sidekiq_daemon/memory_killer_spec.rb

@@ -5,23 +5,14 @@ require 'spec_helper'
 describe Gitlab::SidekiqDaemon::MemoryKiller do
   let(:memory_killer) { described_class.new }
   let(:pid) { 12345 }
-  let(:current_rss_metric) { double('current rss metric') }
-  let(:soft_limit_rss_metric) { double('soft limit rss metric') }
-  let(:hard_limit_rss_metric) { double('hard limit rss metric') }
-  let(:current_phase_metric) { double('current phase metric') }
 
   before do
     allow(Sidekiq.logger).to receive(:info)
     allow(Sidekiq.logger).to receive(:warn)
-    allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_current_rss, anything).and_return(current_rss_metric)
-    allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_memory_killer_soft_limit_rss, anything).and_return(soft_limit_rss_metric)
-    allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_memory_killer_hard_limit_rss, anything).and_return(hard_limit_rss_metric)
-    allow(Gitlab::Metrics).to receive(:gauge).with(:sidekiq_memory_killer_phase, anything).and_return(current_phase_metric)
     allow(memory_killer).to receive(:pid).and_return(pid)
-    allow(current_rss_metric).to receive(:set)
-    allow(soft_limit_rss_metric).to receive(:set)
-    allow(hard_limit_rss_metric).to receive(:set)
-    allow(current_phase_metric).to receive(:set)
+
+    # make sleep no-op
+    allow(memory_killer).to receive(:sleep) {}
   end
 
   describe '#run_thread' do
@@ -51,8 +42,9 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
             pid: pid,
             message: "Exception from run_thread: My Exception")
 
-        expect(memory_killer).to receive(:rss_within_range?).twice.and_raise(StandardError, 'My Exception')
-        expect(memory_killer).to receive(:sleep).twice.with(Gitlab::SidekiqDaemon::MemoryKiller::CHECK_INTERVAL_SECONDS)
+        expect(memory_killer).to receive(:rss_within_range?)
+          .twice
+          .and_raise(StandardError, 'My Exception')
 
         expect { subject }.not_to raise_exception
       end
@@ -64,7 +56,9 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
             pid: pid,
             message: "Exception from run_thread: My Exception")
 
-        expect(memory_killer).to receive(:rss_within_range?).once.and_raise(Exception, 'My Exception')
+        expect(memory_killer).to receive(:rss_within_range?)
+          .once
+          .and_raise(Exception, 'My Exception')
 
         expect(memory_killer).to receive(:sleep).with(Gitlab::SidekiqDaemon::MemoryKiller::CHECK_INTERVAL_SECONDS)
         expect(Sidekiq.logger).to receive(:warn).once
@@ -90,7 +84,9 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
 
     it 'not invoke restart_sidekiq when rss in range' do
-      expect(memory_killer).to receive(:rss_within_range?).twice.and_return(true)
+      expect(memory_killer).to receive(:rss_within_range?)
+        .twice
+        .and_return(true)
 
       expect(memory_killer).not_to receive(:restart_sidekiq)
 
@@ -98,9 +94,12 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
 
     it 'invoke restart_sidekiq when rss not in range' do
-      expect(memory_killer).to receive(:rss_within_range?).at_least(:once).and_return(false)
+      expect(memory_killer).to receive(:rss_within_range?)
+        .at_least(:once)
+        .and_return(false)
 
-      expect(memory_killer).to receive(:restart_sidekiq).at_least(:once)
+      expect(memory_killer).to receive(:restart_sidekiq)
+        .at_least(:once)
 
       subject
     end
@@ -109,10 +108,9 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
   describe '#stop_working' do
     subject { memory_killer.send(:stop_working)}
 
-    it 'changed enable? to false' do
-      expect(memory_killer.send(:enabled?)).to be true
-      subject
-      expect(memory_killer.send(:enabled?)).to be false
+    it 'changes enable? to false' do
+      expect { subject }.to change { memory_killer.send(:enabled?) }
+        .from(true).to(false)
     end
   end
 
@@ -136,7 +134,9 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
       expect(memory_killer).to receive(:get_soft_limit_rss).and_return(200)
       expect(memory_killer).to receive(:get_hard_limit_rss).and_return(300)
 
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:running], 100, 200, 300)
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:running)
+        .and_call_original
 
       expect(Gitlab::Metrics::System).to receive(:monotonic_time).and_call_original
       expect(memory_killer).not_to receive(:log_rss_out_of_range)
@@ -145,11 +145,18 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
 
     it 'return false when rss exceeds hard_limit_rss' do
-      expect(memory_killer).to receive(:get_rss).and_return(400)
+      expect(memory_killer).to receive(:get_rss).at_least(:once).and_return(400)
       expect(memory_killer).to receive(:get_soft_limit_rss).at_least(:once).and_return(200)
       expect(memory_killer).to receive(:get_hard_limit_rss).at_least(:once).and_return(300)
 
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:running], 400, 200, 300)
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:running)
+        .and_call_original
+
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:above_soft_limit)
+        .and_call_original
+
       expect(Gitlab::Metrics::System).to receive(:monotonic_time).and_call_original
 
       expect(memory_killer).to receive(:log_rss_out_of_range).with(400, 300, 200)
@@ -158,15 +165,21 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
 
     it 'return false when rss exceed hard_limit_rss after a while' do
-      expect(memory_killer).to receive(:get_rss).and_return(250, 400)
+      expect(memory_killer).to receive(:get_rss).and_return(250, 400, 400)
       expect(memory_killer).to receive(:get_soft_limit_rss).at_least(:once).and_return(200)
       expect(memory_killer).to receive(:get_hard_limit_rss).at_least(:once).and_return(300)
 
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:running], 250, 200, 300)
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:above_soft_limit], 400, 200, 300)
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:running)
+        .and_call_original
+
+      expect(memory_killer).to receive(:refresh_state)
+        .at_least(:once)
+        .with(:above_soft_limit)
+        .and_call_original
+
       expect(Gitlab::Metrics::System).to receive(:monotonic_time).twice.and_call_original
       expect(memory_killer).to receive(:sleep).with(check_interval_seconds)
-
       expect(memory_killer).to receive(:log_rss_out_of_range).with(400, 300, 200)
 
       expect(subject).to be false
@@ -177,8 +190,14 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
       expect(memory_killer).to receive(:get_soft_limit_rss).and_return(200, 200)
       expect(memory_killer).to receive(:get_hard_limit_rss).and_return(300, 300)
 
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:running], 250, 200, 300)
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:above_soft_limit], 100, 200, 300)
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:running)
+        .and_call_original
+
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:above_soft_limit)
+        .and_call_original
+
       expect(Gitlab::Metrics::System).to receive(:monotonic_time).twice.and_call_original
       expect(memory_killer).to receive(:sleep).with(check_interval_seconds)
 
@@ -187,19 +206,27 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
       expect(subject).to be true
     end
 
-    it 'return false when rss exceed soft_limit_rss longer than GRACE_BALLOON_SECONDS' do
-      expect(memory_killer).to receive(:get_rss).exactly(4).times.and_return(250)
-      expect(memory_killer).to receive(:get_soft_limit_rss).exactly(4).times.and_return(200)
-      expect(memory_killer).to receive(:get_hard_limit_rss).exactly(4).times.and_return(300)
+    context 'when exceeding GRACE_BALLOON_SECONDS' do
+      let(:grace_balloon_seconds) { 0 }
 
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:running], 250, 200, 300)
-      expect(memory_killer).to receive(:update_metrics).exactly(3).times.with(described_class::PHASE[:above_soft_limit], 250, 200, 300)
-      expect(Gitlab::Metrics::System).to receive(:monotonic_time).exactly(5).times.and_call_original
-      expect(memory_killer).to receive(:sleep).exactly(3).times.with(check_interval_seconds).and_call_original
+      it 'return false when rss exceed soft_limit_rss' do
+        allow(memory_killer).to receive(:get_rss).and_return(250)
+        allow(memory_killer).to receive(:get_soft_limit_rss).and_return(200)
+        allow(memory_killer).to receive(:get_hard_limit_rss).and_return(300)
 
-      expect(memory_killer).to receive(:log_rss_out_of_range).with(250, 300, 200)
+        expect(memory_killer).to receive(:refresh_state)
+          .with(:running)
+          .and_call_original
 
-      expect(subject).to be false
+        expect(memory_killer).to receive(:refresh_state)
+          .with(:above_soft_limit)
+          .and_call_original
+
+        expect(memory_killer).to receive(:log_rss_out_of_range)
+          .with(250, 300, 200)
+
+        expect(subject).to be false
+      end
     end
   end
 
@@ -217,19 +244,36 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
 
     it 'send signal' do
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:stop_fetching_new_jobs], 100, 200, 300).ordered
-      expect(memory_killer).to receive(:signal_and_wait).with(shutdown_timeout_seconds, 'SIGTSTP', 'stop fetching new jobs').ordered
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:shutting_down], 100, 200, 300).ordered
-      expect(memory_killer).to receive(:signal_and_wait).with(11, 'SIGTERM', 'gracefully shut down').ordered
-      expect(memory_killer).to receive(:update_metrics).with(described_class::PHASE[:killing_sidekiq], 100, 200, 300).ordered
-      expect(memory_killer).to receive(:signal_pgroup).with('SIGKILL', 'die').ordered
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:stop_fetching_new_jobs)
+        .ordered
+        .and_call_original
+      expect(memory_killer).to receive(:signal_and_wait)
+        .with(shutdown_timeout_seconds, 'SIGTSTP', 'stop fetching new jobs')
+        .ordered
+
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:shutting_down)
+        .ordered
+        .and_call_original
+      expect(memory_killer).to receive(:signal_and_wait)
+        .with(11, 'SIGTERM', 'gracefully shut down')
+        .ordered
+
+      expect(memory_killer).to receive(:refresh_state)
+        .with(:killing_sidekiq)
+        .ordered
+        .and_call_original
+      expect(memory_killer).to receive(:signal_pgroup)
+        .with('SIGKILL', 'die')
+        .ordered
 
       subject
     end
   end
 
   describe '#signal_and_wait' do
-    let(:time) { 7 }
+    let(:time) { 0 }
     let(:signal) { 'my-signal' }
     let(:explanation) { 'my-explanation' }
     let(:check_interval_seconds) { 2 }
@@ -253,14 +297,17 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
 
     it 'send signal and wait till deadline if any job not finished' do
-      expect(Process).to receive(:kill).with(signal, pid).ordered
-      expect(Gitlab::Metrics::System).to receive(:monotonic_time).and_call_original.at_least(:once)
+      expect(Process).to receive(:kill)
+        .with(signal, pid)
+        .ordered
+
+      expect(Gitlab::Metrics::System).to receive(:monotonic_time)
+        .and_call_original
+        .at_least(:once)
 
       expect(memory_killer).to receive(:enabled?).and_return(true).at_least(:once)
       expect(memory_killer).to receive(:any_jobs?).and_return(true).at_least(:once)
 
-      expect(memory_killer).to receive(:sleep).and_call_original.exactly(4).times
-
       subject
     end
   end
@@ -429,14 +476,24 @@ describe Gitlab::SidekiqDaemon::MemoryKiller do
     end
   end
 
-  describe '#update_metrics' do
-    subject { memory_killer.send(:update_metrics, 2, 150, 200, 300) }
+  describe '#refresh_state' do
+    let(:metrics) { memory_killer.instance_variable_get(:@metrics) }
+
+    subject { memory_killer.send(:refresh_state, :shutting_down) }
 
     it 'calls gitlab metrics gauge set methods' do
-      expect(current_phase_metric).to receive(:set).with({}, 2)
-      expect(current_rss_metric).to receive(:set).with({}, 150)
-      expect(soft_limit_rss_metric).to receive(:set).with({}, 200)
-      expect(hard_limit_rss_metric).to receive(:set).with({}, 300)
+      expect(memory_killer).to receive(:get_rss) { 1010 }
+      expect(memory_killer).to receive(:get_soft_limit_rss) { 1020 }
+      expect(memory_killer).to receive(:get_hard_limit_rss) { 1040 }
+
+      expect(metrics[:sidekiq_memory_killer_phase]).to receive(:set)
+        .with({}, described_class::PHASE[:shutting_down])
+      expect(metrics[:sidekiq_current_rss]).to receive(:set)
+        .with({}, 1010)
+      expect(metrics[:sidekiq_memory_killer_soft_limit_rss]).to receive(:set)
+        .with({}, 1020)
+      expect(metrics[:sidekiq_memory_killer_hard_limit_rss]).to receive(:set)
+        .with({}, 1040)
 
       subject
     end