Sidekiq Herkou autoscaler with heroku-api gem

heroku_scaler.rb

require 'heroku-api'

module Background
  class HerokuScaler
    def initialize(
        type = 'worker',
        app = ENV['HEROKU_APP'])
      @client = Heroku::API.new # ENV['HEROKU_API_KEY'] must be present
      @type = type
      @app = app
      @workers = 0
      @known = Time.now - 1
    end

    attr_reader :app
    attr_reader :type

    def workers
      if known?
        @workers
      else
        know client.get_ps(app).body.count {|ps| ps['process'].match /#{type}\.\d?/ }
      end
    end

    def workers=(n)
      if n != @workers || !known?
        p "Scaling #{type} to #{n}"
        client.post_ps_scale(app, type, n)
        know n
      end
    end

    private
    attr_reader :client

    def know(n)
      @known = Time.now + 5
      @workers = n
    end

    def known?
      Time.now < @known
    end
  end
end

sidekiq_autoscaler.rb

require 'securerandom' # bug in Sidekiq as of 2.2.1
require 'sidekiq'

module Background
  module SidekiqAutoscaler
    class Client
      def initialize(scalers)
        @scalers = scalers
      end

      def call(worker_class, item, queue)
        @scalers[queue] && @scalers[queue].workers = 1
        yield
      end
    end

    class Server
      def initialize(scaler, timeout)
        @scaler = scaler
        @timeout = timeout
      end

      def call(worker, msg, queue)
        working!
        yield
      ensure
        working!
        wait_for_task_or_scale
      end

      private
      def registered_queues
        Sidekiq.redis { |x| x.smembers('queues') }
      end

      def empty?(name)
        Sidekiq.redis { |conn| conn.llen("queue:#{name}") == 0 }
      end

      def pending_work?
        registered_queues.any? {|q| !empty?(q)}
      end

      def wait_for_task_or_scale
        loop do
          return if pending_work?
          return @scaler.workers = 0 if idle?
          sleep(0.5)
        end
      end

      def working!
        Sidekiq.redis {|c| c.set('background_activity', Time.now)}
      end

      def idle_time
        Sidekiq.redis {|c|
          t = c.get('background_activity')
          return 0 unless t
          Time.now - Time.parse(t)
        }
      end

      def idle?
        idle_time > @timeout
      end
    end
  end
end

sidekiq_config.rb

require 'securerandom' # bug in Sidekiq as of 2.2.1
require 'sidekiq'
require 'background/sidekiq_autoscaler'
require 'background/heroku_scaler'
redis = ENV['REDIS_URL'] || ENV['REDISTOGO_URL']

# this is a little fancy due to handling multiple process types
heroku = nil
if ENV['HEROKU_APP']
  heroku = {}
  scaleable = %w[default] - (ENV['ALWAYS'] || '').split(' ')
  scaleable.each do |queue|
    heroku[queue] = Background::HerokuScaler.new(
      queue,
      ENV['HEROKU_APP'])
  end
end

Sidekiq.configure_client do |config|
  config.redis = { :size => 1, :url => redis}

  if heroku
    config.client_middleware do |chain|
      chain.add Background::SidekiqAutoscaler::Client, heroku
    end
  end
end

Sidekiq.configure_server do |config|
  config.redis = ConnectionPool.new(:timeout => 4, :size => 4) do
    Redis.connect(:url => redis)
  end

  config.server_middleware do |chain|
    if heroku && ENV['HEROKU_PROCESS'] && heroku[ENV['HEROKU_PROCESS']]
      p "Setting up auto-scaledown"
      chain.add(Background::SidekiqAutoscaler::Server, heroku[ENV['HEROKU_PROCESS']], 60)
    else
      p "Not scaleable"
    end
  end
end