Add 'MemoryKiller' Sidekiq middleware

When enabled, this middleware allows Sidekiq to detect that its RSS has
exceeded a maximum value, triggering a graceful shutdown. This
middleware should be combined with external process supervision that
will restart Sidekiq after the graceful shutdown, such as Runit.

3 files changed, 47 insertions, 1 deletions

diff --git a/CHANGELOG b/CHANGELOG
index 417bd3c2b4d..85178d0dfd7 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -4,7 +4,7 @@ v 7.6.0
   - Add CRON=1 backup setting for quiet backups
   -
   -
-  -
+  - Add optional Sidekiq MemoryKiller middleware (enabled via SIDEKIQ_MAX_RSS env variable)
   -
   -
   -
diff --git a/config/initializers/4_sidekiq.rb b/config/initializers/4_sidekiq.rb
index 228b14cb526..b8a7fd624a5 100644
--- a/config/initializers/4_sidekiq.rb
+++ b/config/initializers/4_sidekiq.rb
@@ -15,6 +15,7 @@ Sidekiq.configure_server do |config|
 
   config.server_middleware do |chain|
     chain.add Gitlab::SidekiqMiddleware::ArgumentsLogger
+    chain.add Gitlab::SidekiqMiddleware::MemoryKiller if ENV['SIDEKIQ_MAX_RSS']
   end
 end
 
diff --git a/lib/gitlab/sidekiq_middleware/memory_killer.rb b/lib/gitlab/sidekiq_middleware/memory_killer.rb
new file mode 100644
index 00000000000..3ef46627916
--- /dev/null
+++ b/lib/gitlab/sidekiq_middleware/memory_killer.rb
@@ -0,0 +1,45 @@
+module Gitlab
+  module SidekiqMiddleware
+    class MemoryKiller
+      # Wait 30 seconds for running jobs to finish during graceful shutdown
+      GRACEFUL_SHUTDOWN_WAIT = 30
+
+      def call(worker, job, queue)
+        yield
+        current_rss = get_rss
+        return unless max_rss > 0 && current_rss > max_rss
+
+        Sidekiq.logger.warn "current RSS #{current_rss} exceeds maximum RSS "\
+         "#{max_rss}"
+        Sidekiq.logger.warn "sending SIGUSR1 to PID #{Process.pid}"
+        Process.kill('SIGUSR1', Process.pid)
+
+        Sidekiq.logger.warn "spawning thread that will send SIGTERM to PID "\
+          "#{Process.pid} in #{graceful_shutdown_wait} seconds"
+        Thread.new do
+          sleep(graceful_shutdown_wait)
+          Process.kill('SIGTERM', Process.pid)
+        end
+      end
+
+      private
+
+      def get_rss
+        output, status = Gitlab::Popen.popen(%W(ps -o rss= -p #{Process.pid}))
+        return 0 unless status.zero?
+
+        output.to_i
+      end
+
+      def max_rss
+        @max_rss ||= ENV['SIDEKIQ_MAX_RSS'].to_s.to_i
+      end
+
+      def graceful_shutdown_wait
+        @graceful_shutdown_wait ||= (
+          ENV['SIDEKIQ_GRACEFUL_SHUTDOWN_WAIT'] || GRACEFUL_SHUTDOWN_WAIT
+        ).to_i
+      end
+    end
+  end
+end