A basic Sidekiq rate limiter

README.md

Don't copypasta this!

Note that this is very purpose-built. The job args in particular depend on a unique key being in the first position.

There is also plenty of brittle code all over the place.

Instead, use this as a guide for how you might build your own. A lot of the parts are the same between implementations.

sidekiq-limiter.rb

class LimitedFetch < Sidekiq::BasicFetch
  attr_reader :job

  def initialize(options)
    super(options)
    @digest = Digest::SHA1.hexdigest(script)
  end

  def retrieve_work
    work = super
    return work if work.blank?

    @job = JSON.parse(work.job) rescue {}
    return work if max_concurrent.blank?

    Sidekiq.redis do |conn|
      if throttled?
        conn.lpush("queue:#{work.queue_name}", work.job)
        nil
      else
        work
      end
    end
  end

  def throttled?
    1 == eval(["limited-throttle-#{id}"], [max_concurrent, 60, @job['jid']])
  end

  def eval(*args)
    Sidekiq.redis { |conn| conn.evalsha(@digest, *args) }
  rescue => e
    raise unless e.message.include?('NOSCRIPT')
    Sidekiq.redis do |conn|
      digest = conn.script(:load, script)
      @digest = digest.freeze

      conn.evalsha(@digest, *args)
    end
  end

  def script
    <<-eos
      local key, limit, ttl, jid = KEYS[1], tonumber(ARGV[1]), tonumber(ARGV[2]), ARGV[3]
      if limit <= redis.call("SCARD", key) and 0 == redis.call("SISMEMBER", key, jid) then
        return 1
      end
      redis.call("SADD", key, jid)
      redis.call("EXPIRE", key, ttl)
      return 0
    eos
  end

  def max_concurrent
    job['max_concurrent']
  end

  def id
    job['args'].first
  end
end

class LimitedMiddleware
  def call(_worker, job, _queue)
    yield
  ensure
    Sidekiq.redis { |conn| conn.srem("limited-throttle-#{job['args'].first}", job['jid']) }
  end
end

Oh, duh, you need to install it too:

Sidekiq.configure_server do |config|
  config.server_middleware do |chain|
    chain.add LimitedMiddleware
  end
  Sidekiq.options[:fetch] = LimitedFetch
end