Sidekiq/Heroku MemoryKiller middleware

memory_killer.rb

# frozen_string_literal: true

require 'platform-api'

# Adapted from GitLab
#
# Docs:   https://docs.gitlab.com/ee/administration/operations/sidekiq_memory_killer.html
# Source: https://gitlab.com/gitlab-org/gitlab-foss/-/blob/master/lib/gitlab/sidekiq_daemon/memory_killer.rb
class Sidekiq::Middleware::MemoryKiller
  # Default the RSS limit to 0, meaning the MemoryKiller is disabled (kilobytes)
  MAX_RSS = Integer(ENV['SIDEKIQ_MEMORY_KILLER_MAX_RSS'] || 0)
  # Give Sidekiq 2 minutes of grace time after exceeding the RSS limit
  GRACE_TIME = Integer(ENV['SIDEKIQ_MEMORY_KILLER_GRACE_TIME_IN_SECONDS'] || 120)

  # Heroku application details
  API_TOKEN = ENV['HEROKU_API_TOKEN']
  APP_NAME = ENV['HEROKU_APP_NAME']
  DYNO = ENV['DYNO']

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

  def call(worker, job, queue) # rubocop:disable Lint/UnusedMethodArgument
    yield
    current_rss = fetch_rss

    return if current_rss < MAX_RSS

    Thread.new do
      # Make sure another thread isn't already waiting to shut Sidekiq down
      if MUTEX.try_lock
        Sidekiq.logger.warn "Current RSS #{current_rss} exceeds maximum RSS #{MAX_RSS}!"
        Sidekiq.logger.warn "This dyno (#{DYNO}) will restart in #{GRACE_TIME} seconds"
        sleep(GRACE_TIME)

        Sidekiq.logger.warn "Restarting dyno (#{DYNO})..."
        heroku = PlatformAPI.connect_oauth(API_TOKEN)
        heroku.dyno.restart(APP_NAME, DYNO)
      end
    end
  end

  private

  def fetch_rss
    Integer(`ps -o rss= -p #{::Process.pid}`)
  end
end

sidekiq.rb

Sidekiq.configure_server do |config|
  config.server_middleware do |chain|
    # Automatically restart Heroku dynos when memory increases.
    if Rails.env.staging? || Rails.env.production?
      chain.add Sidekiq::Middleware::MemoryKiller
    end
  end
end