Skip to content

Instantly share code, notes, and snippets.

@somesh-m
Last active November 16, 2023 05:17
Show Gist options
  • Save somesh-m/5327ce391ec71cbcea2fc8927088e256 to your computer and use it in GitHub Desktop.
Save somesh-m/5327ce391ec71cbcea2fc8927088e256 to your computer and use it in GitHub Desktop.
Lua script used for atomically handling api rate counter
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