MemoryUsageMiddleware

sidekiq_memory_usage_middleware.rb

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

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

  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_#{MemoryUsageMiddleware.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
    MemoryUsageMiddleware.mb
  end
end

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

sidekiq_memory_usage_middleware_spec.rb

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

describe MemoryUsageMiddleware do
  let(:process_mem) { double(:mb => BigDecimal('100')) }
  let(:logger)      { double(:info => true, :error => true) }
  let(:mb)          { BigDecimal('60.0') }
  let(:middleware)  { MemoryUsageMiddleware.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( MemoryUsageMiddleware.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)

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

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

      MemoryUsageMiddleware.mb.should == BigDecimal('120')
      MemoryUsageMiddleware.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")

      MemoryUsageMiddleware.mb
    end
  end

  describe "#call" do
    it "should call MemoryUsageMiddleware.mb" do
      expect( MemoryUsageMiddleware ).to receive(:mb)

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