Last active
November 16, 2023 05:17
-
-
Save somesh-m/5327ce391ec71cbcea2fc8927088e256 to your computer and use it in GitHub Desktop.
Lua script used for atomically handling api rate counter
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
local logchannel = "tokenbucketdebug" | |
-- This function will be used to log debug messages to redis pub/sub channel | |
local function logit(msg) | |
redis.pcall("PUBLISH", logchannel, msg) | |
end | |
local key = KEYS[1] | |
local logPrependText = "User: "..key | |
local availableToken = redis.call('GET', 'tokens_'..key) | |
if availableToken == false then | |
availableToken = 0 | |
end | |
availableToken = tonumber(availableToken) | |
if availableToken > 0 then | |
-- Since we have valid tokens, just decrease the token count and return true | |
logit(logPrependText.." Token Count "..availableToken) | |
redis.call('decr', 'tokens_'..key) | |
return 'True' | |
else | |
-- If token is not present, check if the key with TTL exists, if it doesn' refil the token bucket | |
-- else return false | |
local refilTime = redis.call('GET', 'refill_'..key) | |
if type(refilTime) == 'boolean' and refilTime == false then | |
logit(logPrependText.." Refilling token bucket") | |
redis.call('SET', 'refill_'..key, key) | |
redis.call('EXPIRE', 'refill_'..key, 60) | |
redis.call('SET', 'tokens_'..key, 5) | |
return 'True' | |
else | |
logit(logPrependText.."Time limit hasn't expired") | |
return 'False' | |
end | |
end |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment