Skip to content

Instantly share code, notes, and snippets.

@hschne
Created January 11, 2024 17:02
Show Gist options
  • Save hschne/ae9b3e8da57e96f00ccfba7d8dad1cd9 to your computer and use it in GitHub Desktop.
Save hschne/ae9b3e8da57e96f00ccfba7d8dad1cd9 to your computer and use it in GitHub Desktop.
Leaky Bucket Rate Limiter in Ruby
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
rate_limiter = RateLimit.new("send_mail:#{receiver.id}", redis)
# Raises RateLimit::Error if rate limit exceeded
rate_limiter.limit { send_mail(receiver) }
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment