Leaky Bucket Rate Limiter in Ruby

rate_limit.rb

 frozen_string_literal: true

# A leaky bucket rate limiter for Ruby
#
# @see https://www.mikeperham.com/2020/11/09/the-leaky-bucket-rate-limiter/
# @see https://en.wikipedia.org/wiki/Leaky_bucket
class RateLimit
  class Error < StandardError
    attr_accessor :retry_in

    def initialize(retry_in)
      self.retry_in = retry_in
      super("Rate limit exceeded, please try again in #{retry_in} seconds.")
    end
  end

  # Create a new rate limit.
  #
  # The rate limiter uses Redis under the hood to store the current load and the time of the last request.
  #
  # @param [String] key the key to use for this rate limit. This should be unique to the rate limit you want to create.
  # @option [Redis] redis the redis instance to use.
  # @option [Integer] threshold the number of requests allowed in the interval. Defaults to 10
  # @option [Integer] interval the interval in seconds. Defaults to 60 (1 minute)
  def initialize(key, redis, threshold: 10, interval: 60, redis: RedisInstances.secondary)
    @key = "rate_limit:#{key}"
    @threshold = threshold
    @interval = interval
    @redis = redis
  end

  # Rate limit a certain action in a block
  #
  # @yield the action to execute with a rate limit. This is optional
  # @raise [RateLimit::Error] if the rate limit is exceeded
  def limit(&block)
    bucket = @redis.hgetall(@key).symbolize_keys.transform_values(&:to_i)
    bucket = create_new_counter if bucket.empty?
    leak(bucket)

    if @redis.hget(@key, :current_load).to_i >= @threshold
      retry_in = (@interval / @threshold) - (Time.now.to_i - bucket[:last_request_made_at])
      raise(RateLimit::Error, retry_in)
    end

    increment
    yield(block) if block
  end

  private

  def leak(bucket)
    now = Time.now.to_i
    # Leak rate is based on the interval and the threshold
    leak_rate = @interval / @threshold
    leak_amount = (now - bucket[:last_request_made_at]) / leak_rate
    # Decrement the bucket load, or empty it if it's been a long enough time for it to leak everything
    bucket_load = [bucket[:current_load] - leak_amount, 0].max
    @redis.hset(@key, :current_load, bucket_load)
  end

  def increment
    @redis.multi do |multi|
      multi.hincrby(@key, :current_load, 1)
      multi.hset(@key, :last_request_made_at, Time.now.to_i)
    end
  end

  def create_new_counter
    @redis.hset(@key, :current_load, 0, :last_request_made_at, 0)
    @redis.expire(@key, 604800) # A bucket expires after one week
    { current_load: 0, last_request_made_at: 0 }
  end
end

usage.rb

rate_limiter = RateLimit.new("send_mail:#{receiver.id}", redis)
# Raises RateLimit::Error if rate limit exceeded
rate_limiter.limit { send_mail(receiver) }