Created
November 24, 2022 15:17
-
-
Save fpaint/e51a7bd09dd9c051472997df3c4fca62 to your computer and use it in GitHub Desktop.
Rate limiter class
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| # frozen_string_literal: true | |
| # Redis-based rate limiter inspired by a leaking bucket algorithm | |
| class RateLimiter | |
| BUCKET_SIZE = 1000 # arbitrary number, for normalization | |
| BUCKET_MAXIMUM = 5000 # upper limiter of score | |
| TTL = 3600 # expiration time of used keys, to keep the redis clean | |
| OK_RESPONSE = "OK" | |
| COOLDOWN_RESPONSE = "COOLDOWN" | |
| OVERFLOW_RESPONSE = "OVERFLOW" | |
| # key - identifier of request to be limited, "mutation_name + device_uid" for example | |
| # cooldown - min time between requests, in seconds, can be float | |
| # period - time window, in seconds | |
| # limit - allowed amount of requests per period (0 to disable rate limit) | |
| # current_time - timestamp | |
| def self.call(key, cooldown, period, limit, current_time = nil) | |
| current_time ||= Time.now.to_f | |
| $redis.evalsha(script_sha, ["rate_limiter:#{key}"], [cooldown, period, limit, current_time]) | |
| end | |
| # for testing purposes | |
| def self.preset(key, gauge, current_time = nil) | |
| current_time ||= Time.now.to_f | |
| $redis.call("set", "rate_limiter:#{key}:last_call", current_time, "EX", TTL) | |
| $redis.call("set", "rate_limiter:#{key}:gauge", gauge, "EX", TTL) | |
| end | |
| # also | |
| def self.unset(key) | |
| $redis.del("rate_limiter:#{key}:last_call", "rate_limiter:#{key}:gauge") | |
| end | |
| def self.script_sha | |
| @script_sha ||= $redis.script("LOAD", script_source) | |
| end | |
| def self.script_source | |
| %{ | |
| local key = KEYS[1] | |
| local last_call_key = key..':last_call' | |
| local gauge_key = key..':gauge' | |
| local cooldown = tonumber(ARGV[1]) | |
| local period = tonumber(ARGV[2]) | |
| local limit = tonumber(ARGV[3]) | |
| local current_time = tonumber(ARGV[4]) | |
| local request_price = limit > 0 and (#{BUCKET_SIZE} / limit) or 0.0 | |
| local drops_per_second = period / #{BUCKET_SIZE} | |
| local last_call = redis.call('GETSET', last_call_key, current_time) | |
| redis.call('EXPIRE', last_call_key, #{TTL}) | |
| if last_call == false then | |
| redis.call('SET', gauge_key, request_price, 'EX', #{TTL}) | |
| return "#{OK_RESPONSE}" | |
| end | |
| local delta = current_time - tonumber(last_call) | |
| local gauge = tonumber(redis.call('GET', gauge_key)) or 0.0 | |
| gauge = math.min(#{BUCKET_MAXIMUM}, math.max(0.0, gauge - delta * drops_per_second + request_price)) | |
| redis.call('SET', gauge_key, tostring(gauge), 'EX', #{TTL}) | |
| if delta < cooldown then | |
| return "#{COOLDOWN_RESPONSE}" | |
| end | |
| if gauge > #{BUCKET_SIZE} then | |
| return "#{OVERFLOW_RESPONSE}" | |
| end | |
| return "#{OK_RESPONSE}" | |
| } | |
| end | |
| end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment