Skip to content

Commit

Permalink
Add TTL to Redis token bucket implementation (#1757)
Browse files Browse the repository at this point in the history
  • Loading branch information
pilcrowonpaper authored Dec 9, 2024
1 parent 5c50a0a commit c39c1bf
Showing 1 changed file with 16 additions and 7 deletions.
23 changes: 16 additions & 7 deletions pages/rate-limit/token-bucket.md
Original file line number Diff line number Diff line change
Expand Up @@ -73,27 +73,36 @@ local cost = tonumber(ARGV[3])
local now = tonumber(ARGV[4]) -- Current unix time in seconds

local fields = redis.call("HGETALL", key)

if #fields == 0 then
redis.call("HSET", key, "count", max - cost, "refilled_at", now)
return {1}
local expiresInSeconds = cost * refillIntervalSeconds
redis.call("HSET", key, "count", max - cost, "refilled_at", now)
redis.call("EXPIRE", key, expiresInSeconds)
return {1}
end

local count = 0
local refilledAt = 0
for i = 1, #fields, 2 do
if fields[i] == "count" then
count = tonumber(fields[i+1])
elseif fields[i] == "refilled_at" then
refilledAt = tonumber(fields[i+1])
end
count = tonumber(fields[i+1])
elseif fields[i] == "refilled_at" then
refilledAt = tonumber(fields[i+1])
end
end

local refill = math.floor((now - refilledAt) / refillIntervalSeconds)
count = math.min(count + refill, max)
refilledAt = refilledAt + refill * refillIntervalSeconds

if count < cost then
return {0}
return {0}
end

count = count - cost
local expiresInSeconds = (max - count) * refillIntervalSeconds
redis.call("HSET", key, "count", count, "refilled_at", now)
redis.call("EXPIRE", key, expiresInSeconds)
return {1}
```

Expand Down

0 comments on commit c39c1bf

Please sign in to comment.