Heroku Autoscale middleware for Rails

heroku-autoscale.rb

# Based ON https://github.com/ddollar/heroku-autoscale/blob/master/lib/heroku/autoscale.rb
# This middleware implements a simple dyno scaler for heroku + rails. Requires Redis.

# in Gemfile
# gem 'redlock'

# in an initializer
REDLOCK = Redlock::Client.new([ENV['REDIS_URL']], { retry_count: 0 })
REDIS = RedisClient.new(url: ENV['REDIS_URL'])

# in lib/middleware
module Heroku
  class Autoscale
    attr_reader :app, :options

    LOCK_KEY = "heroku-autoscale:lock"
    AVG_KEY = "heroku-autoscale:rolling_average"

    def initialize(app, options = {})
      @app = app
      @options = default_options.merge(options)
      check_options!
    end

    def call(env)
      # Autoscale only triggers after the existing lock expires
      Thread.new { autoscale(env) } if REDLOCK.lock(LOCK_KEY, options.fetch(:min_frequency))

      app.call(env)
    end

    private ######################################################################

    def autoscale(env)
      original_dynos = dynos = current_dynos
      wait = queue_wait(env)
      average_wait = average_wait(wait)
      Rails.logger.info "Heroku::Autoscale: wait time is #{wait}ms avg wait time is #{average_wait}ms"

      dynos -= 1 if average_wait <= options.fetch(:queue_wait_low)
      dynos += 1 if average_wait >= options.fetch(:queue_wait_high)

      dynos = options.fetch(:min_dynos) if dynos < options.fetch(:min_dynos)
      dynos = options.fetch(:max_dynos) if dynos > options.fetch(:max_dynos)
      dynos = 1 if dynos < 1

      Rails.logger.info "Heroku::Autoscale: current_dynos=#{original_dynos} new_dynos=#{dynos}"

      if dynos != original_dynos
        Rails.logger.info "Heroku::Autoscale: scaling from #{original_dynos} to #{dynos}"
        set_dynos(dynos)
      end
    end

    def average_wait(time)
      max_samples = 60000 / 10000 * 2 # samples in 2 minutes

      REDIS.multi do
        REDIS.call('lpush', AVG_KEY, time)
        REDIS.call('ltrim', AVG_KEY, 0, max_samples)
      end
      REDIS.call('lrange', AVG_KEY, 0, -1).sum(&:to_i) / REDIS.call('llen', AVG_KEY).to_f
    end

    def check_options!
      errors = []
      errors << "Must supply :heroku_api_key to Heroku::Autoscale" unless options.fetch(:heroku_api_key)
      errors << "Must supply :app_name to Heroku::Autoscale" unless options.fetch(:app_name)
      raise errors.join(" / ") unless errors.empty?
    end

    def current_dynos
      heroku.formation.info(options.fetch(:app_name), 'web')['quantity']
    end

    def default_options
      {
        min_dynos: 1,
        max_dynos: 10,
        queue_wait_low: 200,
        queue_wait_high: 300,
        min_frequency: 15000,
        dyno_size: 'standard-1x',
      }
    end

    def heroku
      @heroku ||= PlatformAPI.connect_oauth(options.fetch(:heroku_api_key))
    end

    def queue_wait(env)
      request_start = Time.at(env["HTTP_X_REQUEST_START"].to_i / 1000.0).getlocal
      puma_network_time = env["puma.request_body_wait"].to_f # milliseconds

      ((Time.current - request_start) * 1000) - puma_network_time
    end

    def set_dynos(count)
      heroku.formation.update(
        options.fetch(:app_name),
        'web',
        quantity: count,
        size: options.fetch(:dyno_size),
      )
    end
  end
end

# in production.rb
  config.middleware.use Heroku::Autoscale,
    app_name: 'application-name',
    heroku_api_key: ENV.fetch('HEROKU_API_KEY'),
    min_dynos: 1,
    max_dynos: 10,
    queue_wait_low: 200,
    queue_wait_high: 300,
    min_frequency: 15000,
    dyno_size: 'standard-1x'