Sidekiq Max Memory Middleware

sidekiq_max_memory_middleware.rb

# config/initializers/sidekiq_max_memory_middleware.rb
require 'socket'

class MaxMemoryMiddleware
  attr_reader :active, :max_oom, :shutdown
  cattr_reader :cached_mb, :last_checked_at, :hostname

  def initialize
    if ENV['MAX_SIDEKIQ_OOM']
      @active   = true
      @shutdown = false
      @max_oom  = BigDecimal(ENV['MAX_SIDEKIQ_OOM'])
    else
      @active = false
    end
  end

  def self.start
    @@start ||= Time.now
  end

  def self.hostname
    @@hostname ||= Socket.gethostname.split('.').first
  end

  def self.mb
    time = Time.now
    if @@cached_mb.nil? || (@@last_checked_at && (time - @@last_checked_at > 30.0))
      @@last_checked_at = time
      GetProcessMem.new.mb.tap do |mb|
        Appsignal.set_gauge("sidekiq_mem_#{MaxMemoryMiddleware.hostname}_#{Process.pid}", mb.round)
        Sidekiq.logger.info "Memory usage #{mb} MB"
        @@cached_mb = mb
      end
    else
      @@cached_mb
    end
  end

  def call(worker, job, queue)
    yield
  ensure
    if @active &&
      !@shutdown &&
      MaxMemoryMiddleware.mb > @max_oom &&
      MaxMemoryMiddleware.start < 10.minutes.ago
      @shutdown = true
      Appsignal.increment_counter('sidekiq_oom', 1)
      Sidekiq.logger.error "Exceeded max memory limit (#{MaxMemoryMiddleware.mb} > #{@max_oom} MB)"
      Process.kill('TERM', Process.pid)
    end
  end
end

Sidekiq.configure_server do |config|
  config.server_middleware do |chain|
    chain.add MaxMemoryMiddleware
  end
end

sidekiq_max_memory_middleware_spec.rb

# spec/initializers/sidekiq_max_memory_middleware_spec.rb
require 'spec_helper'

describe MaxMemoryMiddleware do
  let(:process_mem) { double(:mb => BigDecimal('100')) }
  let(:logger)      { double(:info => true, :error => true) }
  let(:mb)          { BigDecimal('60.0') }
  let(:middleware)  { MaxMemoryMiddleware.new }

  before do
    Process.stub(:pid => 123)
    Socket.stub(:gethostname => 'worker1.staging')
    GetProcessMem.stub(:new => process_mem)
    Sidekiq.stub(:logger => logger)
    Process.stub(:kill => true)
  end

  describe ".hostname" do
    it "should return the first part of the hostname, without `.`" do
      expect( MaxMemoryMiddleware.hostname ).to eql('worker1')
    end
  end

  describe ".mb" do
    after do
      Timecop.return
    end

    it "should return the mb and cache it for 30 seconds" do
      process_mem.
        should_receive(:mb).
        twice.
        and_return(BigDecimal('100'), BigDecimal('120'))

      time = Time.utc(2014, 1, 1, 0, 0)
      Timecop.freeze Time.utc(2014, 1, 1, 0, 0)

      MaxMemoryMiddleware.mb.should == BigDecimal('100')
      MaxMemoryMiddleware.mb.should == BigDecimal('100')
      MaxMemoryMiddleware.mb.should == BigDecimal('100')
      MaxMemoryMiddleware.mb.should == BigDecimal('100')
      MaxMemoryMiddleware.last_checked_at.should == time

      new_time = Time.utc(2014, 1, 1, 0, 2)
      Timecop.freeze new_time

      MaxMemoryMiddleware.mb.should == BigDecimal('120')
      MaxMemoryMiddleware.last_checked_at.should == new_time
    end

    it "should log the memory usage" do
      expect( Appsignal ).to receive(:set_gauge)
        .with("sidekiq_mem_worker1_123", 100)

      expect( logger ).to receive(:info)
        .with("Memory usage 100.0 MB")

      MaxMemoryMiddleware.mb
    end
  end

  context "when not active" do
    before { ENV.delete('MAX_SIDEKIQ_OOM') }

    its(:active) { should be_false }

    it "should yield and do nothing when called" do
      expect { |b|
        subject.call(nil, nil, nil, &b)
      }.to yield_with_no_args
    end
  end

  context "with min and max set" do
    before do
      ENV['MAX_SIDEKIQ_OOM'] = '100'
      MaxMemoryMiddleware.stub(:mb => mb)
    end

    its(:active)  { should be_true }
    its(:max_oom) { should == BigDecimal('100.0') }

    context "below limit" do
      let(:mb) { BigDecimal('60.0') }

      it "should not exit if below limit" do
        expect( Process ).to_not receive(:kill)
      end
    end

    context "above limit" do
      let(:mb) { BigDecimal('120.0') }
      before { MaxMemoryMiddleware.stub(:start => 11.minutes.ago) }

      it "should exit if above limit" do
        expect( Process ).to receive(:kill).with('TERM', 123)
      end

      it "should log the shutdown" do
        expect( logger ).to receive(:error)
          .with("Exceeded max memory limit (120.0 > 100.0 MB)")
      end

      it "should send metrics to AppSignal" do
        expect( Appsignal ).to receive(:increment_counter)
          .with('sidekiq_oom', 1)
      end

      context "when shutting down" do
        before { middleware.call(nil, nil, nil) { } }

        it "should not kill the process again" do
          expect( Process ).to_not receive(:kill)
        end
      end

      context "when in the cooldown phase" do
        before { MaxMemoryMiddleware.stub(:start => 8.minutes.ago) }

        it "should not kill the process again" do
          expect( Process ).to_not receive(:kill)
        end
      end
    end

    after { middleware.call(nil, nil, nil) { } }
  end
end