-
Notifications
You must be signed in to change notification settings - Fork 65
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Rate limiter: add an alternative implementation based on Redis. (#620)
* Add a Redis rate limiter * Add useful comments * Allow Redis rate limiter to be used * Regenerate readme * Fix warning * Fix tests in machine in CircleCI * Tidy up go modules Signed-off-by: Douglas Camata <159076+douglascamata@users.noreply.github.com> * Address PR review comments * Regenerate readme * Fix test assertion to account for negative time delta * Fix Redis limiter tests * Adjust test case comment * Fix/improve code comments * Keep Go 1.20 * Tidy with go1.20 Signed-off-by: Douglas Camata <159076+douglascamata@users.noreply.github.com> --------- Signed-off-by: Douglas Camata <159076+douglascamata@users.noreply.github.com>
- Loading branch information
1 parent
1132fdb
commit f2995f3
Showing
9 changed files
with
369 additions
and
14 deletions.
There are no files selected for viewing
This file contains 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
This file contains 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
This file contains 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
This file contains 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
This file contains 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
This file contains 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
This file contains 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,66 @@ | ||
| -- this script has side-effects, so it requires replicate commands mode | ||
| redis.replicate_commands() | ||
|
|
||
| local rate_limit_key = KEYS[1] -- The key to the rate limit bucket. | ||
| local now = ARGV[1] -- Current time (Unix time in milliseconds). | ||
| local burst = ARGV[2] -- This represents the total capacity of the bucket. | ||
| local rate = ARGV[3] -- This represents the amount that leaks from the bucket. | ||
| local period = ARGV[4] -- This represents how often the "rate" leaks from the bucket (in milliseconds). | ||
| local cost = ARGV[5] -- This represents the cost of the request. Often 1 is used per request. | ||
| -- It allows some requests to be assigned a higher cost. | ||
|
|
||
| local emission_interval = period / rate | ||
| local increment = emission_interval * cost | ||
| local burst_offset = emission_interval * burst | ||
|
|
||
| local tat = redis.call("GET", rate_limit_key) | ||
|
|
||
| if not tat then | ||
| tat = now | ||
| else | ||
| tat = tonumber(tat) | ||
| end | ||
| tat = math.max(tat, now) | ||
|
|
||
| local new_tat = tat + increment | ||
| local allow_at = new_tat - burst_offset | ||
| local diff = now - allow_at | ||
|
|
||
| local limited | ||
| local retry_in | ||
| local reset_in | ||
|
|
||
| local remaining = math.floor(diff / emission_interval) -- poor man's round | ||
|
|
||
| if remaining < 0 then | ||
| limited = 1 | ||
| -- calculate how many tokens there actually are, since | ||
| -- remaining is how many there would have been if we had been able to limit | ||
| -- and we did not limit | ||
| remaining = math.floor((now - (tat - burst_offset)) / emission_interval) | ||
| reset_in = math.ceil(tat - now) | ||
| retry_in = math.ceil(diff * -1) | ||
| elseif remaining == 0 and increment <= 0 then | ||
| -- request with cost of 0 | ||
| -- cost of 0 with remaining 0 is still limited | ||
| limited = 1 | ||
| remaining = 0 | ||
| reset_in = math.ceil(tat - now) | ||
| retry_in = 0 -- retry in is meaningless when cost is 0 | ||
| else | ||
| limited = 0 | ||
| reset_in = math.ceil(new_tat - now) | ||
| retry_in = 0 | ||
| if increment > 0 then | ||
| redis.call("SET", rate_limit_key, new_tat, "PX", reset_in) | ||
| end | ||
| end | ||
|
|
||
| -- return values (in order): | ||
| -- limited = integer-encoded boolean, 1 if limited, 0 if not | ||
| -- remaining = number of tokens remaining | ||
| -- retry_in = milliseconds until the next request will be allowed | ||
| -- reset_in = milliseconds until the rate limit window resets | ||
| -- diff = milliseconds since the last request | ||
| -- emission_interval = milliseconds between token emissions | ||
| return {limited, remaining, retry_in, reset_in, tostring(diff), tostring(emission_interval)} |
This file contains 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,82 @@ | ||
| package ratelimit | ||
|
|
||
| import ( | ||
| "context" | ||
| _ "embed" | ||
| "strconv" | ||
| "time" | ||
|
|
||
| "github.com/redis/rueidis" | ||
| ) | ||
|
|
||
| //go:embed gcra_rate_limit.lua | ||
| var gcraRateLimitScript string | ||
|
|
||
| // RedisRateLimiter is a type that represents a rate limiter that uses Redis as its backend. | ||
| // The rate limiting is a leaky bucket implementation using the generic cell rate algorithm. | ||
| // See https://en.wikipedia.org/wiki/Generic_cell_rate_algorithm for details on how this algorithm works. | ||
| type RedisRateLimiter struct { | ||
| client rueidis.Client | ||
| } | ||
|
|
||
| // Ensure RedisRateLimiter implements the SharedRateLimiter interface. | ||
| var _ SharedRateLimiter = (*RedisRateLimiter)(nil) | ||
|
|
||
| // NewRedisRateLimiter creates a new instance of RedisRateLimiter. | ||
| func NewRedisRateLimiter(addresses []string) (*RedisRateLimiter, error) { | ||
| client, err := rueidis.NewClient(rueidis.ClientOption{InitAddress: addresses}) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| return &RedisRateLimiter{client: client}, nil | ||
| } | ||
|
|
||
| // GetRateLimits retrieves the rate limits for a given request using a Redis Rate Limiter. | ||
| // It returns the amount of remaining requests, the reset time in milliseconds, and any error that occurred. | ||
| func (r *RedisRateLimiter) GetRateLimits(ctx context.Context, req *request) (remaining, resetTime int64, err error) { | ||
| inspectScript := rueidis.NewLuaScript(gcraRateLimitScript) | ||
| rateLimitParameters := []string{ | ||
| strconv.FormatInt(time.Now().UnixMilli(), 10), // now | ||
| strconv.FormatInt(req.limit, 10), // burst | ||
| strconv.FormatInt(req.limit, 10), // rate | ||
| strconv.FormatInt(req.duration, 10), // period | ||
| "1", // cost | ||
| } | ||
| result := inspectScript.Exec(ctx, r.client, []string{req.key}, rateLimitParameters) | ||
| limited, remaining, resetIn, err := r.parseRateLimitResult(&result) | ||
| if err != nil { | ||
| return 0, 0, err | ||
| } | ||
| resetTime = time.Now().Add(time.Duration(resetIn) * time.Millisecond).UnixMilli() | ||
| if limited { | ||
| return remaining, resetTime, errOverLimit | ||
| } | ||
| return remaining, resetTime, nil | ||
| } | ||
|
|
||
| // parseRateLimitResult parses the result of a rate limit check from Redis. | ||
| // It takes a RedisResult as input and returns the parsed rate limit values: whether the request is limited, | ||
| // the number of remaining requests, the reset time in milliseconds, and any error that occurred during parsing. | ||
| func (r *RedisRateLimiter) parseRateLimitResult(result *rueidis.RedisResult) (limited bool, remaining, resetIn int64, err error) { | ||
| values, err := result.ToArray() | ||
| if err != nil { | ||
| return false, 0, 0, err | ||
| } | ||
|
|
||
| limited, err = values[0].AsBool() | ||
| if err != nil { | ||
| return false, 0, 0, err | ||
| } | ||
|
|
||
| remaining, err = values[1].AsInt64() | ||
| if err != nil { | ||
| return false, 0, 0, err | ||
| } | ||
|
|
||
| resetIn, err = values[3].AsInt64() | ||
| if err != nil { | ||
| return false, 0, 0, err | ||
| } | ||
|
|
||
| return limited, remaining, resetIn, nil | ||
| } |
Oops, something went wrong.