Move rate limiter logic to Lua scripts

This commit is contained in:
Jon Chambers
2020-07-06 10:10:13 -04:00
committed by GitHub
parent f5ddb0f1f8
commit b585c6676d
8 changed files with 214 additions and 294 deletions

View File

@@ -0,0 +1,33 @@
local bucketId = KEYS[1]
local bucketSize = tonumber(ARGV[1])
local leakRatePerMillis = tonumber(ARGV[2])
local currentTimeMillis = tonumber(ARGV[3])
local amount = tonumber(ARGV[4])
local leakyBucket
if redis.call("EXISTS", bucketId) == 1 then
leakyBucket = cjson.decode(redis.call("GET", bucketId))
else
leakyBucket = {
bucketSize = bucketSize,
leakRatePerMillis = leakRatePerMillis,
spaceRemaining = bucketSize,
lastUpdateTimeMillis = currentTimeMillis
}
end
local elapsedTime = currentTimeMillis - leakyBucket["lastUpdateTimeMillis"]
local updatedSpaceRemaining = math.min(leakyBucket["bucketSize"], math.floor(leakyBucket["spaceRemaining"] + (elapsedTime * leakyBucket["leakRatePerMillis"])))
redis.call("SET", "elapsedTime", elapsedTime)
redis.call("SET", "updatedSpaceRemaining", updatedSpaceRemaining)
if updatedSpaceRemaining >= amount then
leakyBucket["spaceRemaining"] = updatedSpaceRemaining - amount
redis.call("SET", bucketId, cjson.encode(leakyBucket))
return true
else
return false
end