BigW Consortium Gitlab

memory_killer.rb 2.05 KB
Newer Older
1 2 3
module Gitlab
  module SidekiqMiddleware
    class MemoryKiller
4 5
      # Default the RSS limit to 0, meaning the MemoryKiller is disabled
      MAX_RSS = (ENV['SIDEKIQ_MEMORY_KILLER_MAX_RSS'] || 0).to_s.to_i
6
      # Give Sidekiq 15 minutes of grace time after exceeding the RSS limit
7
      GRACE_TIME = (ENV['SIDEKIQ_MEMORY_KILLER_GRACE_TIME'] || 15 * 60).to_s.to_i
8
      # Wait 30 seconds for running jobs to finish during graceful shutdown
9
      SHUTDOWN_WAIT = (ENV['SIDEKIQ_MEMORY_KILLER_SHUTDOWN_WAIT'] || 30).to_s.to_i
10
      SHUTDOWN_SIGNAL = (ENV['SIDEKIQ_MEMORY_KILLER_SHUTDOWN_SIGNAL'] || 'SIGKILL').to_s
11 12 13

      # Create a mutex used to ensure there will be only one thread waiting to
      # shut Sidekiq down
14
      MUTEX = Mutex.new
15 16 17 18

      def call(worker, job, queue)
        yield
        current_rss = get_rss
19

20
        return unless MAX_RSS > 0 && current_rss > MAX_RSS
21

Jacob Vosmaer committed
22
        Thread.new do
23 24 25 26
          # Return if another thread is already waiting to shut Sidekiq down
          return unless MUTEX.try_lock

          Sidekiq.logger.warn "current RSS #{current_rss} exceeds maximum RSS "\
27
            "#{MAX_RSS}"
28
          Sidekiq.logger.warn "this thread will shut down PID #{Process.pid} - Worker #{worker.class} - JID-#{job['jid']}"\
29
            "in #{GRACE_TIME} seconds"
30
          sleep(GRACE_TIME)
31

James Lopez committed
32
          Sidekiq.logger.warn "sending SIGTERM to PID #{Process.pid} - Worker #{worker.class} - JID-#{job['jid']}"
33
          Process.kill('SIGTERM', Process.pid)
34

35
          Sidekiq.logger.warn "waiting #{SHUTDOWN_WAIT} seconds before sending "\
James Lopez committed
36
            "#{SHUTDOWN_SIGNAL} to PID #{Process.pid} - Worker #{worker.class} - JID-#{job['jid']}"
37
          sleep(SHUTDOWN_WAIT)
38

39
          Sidekiq.logger.warn "sending #{SHUTDOWN_SIGNAL} to PID #{Process.pid} - Worker #{worker.class} - JID-#{job['jid']}"
40
          Process.kill(SHUTDOWN_SIGNAL, Process.pid)
41 42 43 44 45 46 47 48 49 50 51 52 53 54
        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
    end
  end
end