Skip to content

Commit

Permalink
[Minor] Improve style in the scripts
Browse files Browse the repository at this point in the history
  • Loading branch information
@vstakhov
vstakhov committed Mar 25, 2023
1 parent 0dd7203 commit 5d8aa83
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 40 deletions.
33 changes: 18 additions & 15 deletions lualib/redis_scripts/ratelimit_check.lua
View file
Expand Up @@ -16,47 +16,50 @@
local last = redis.call('HGET', KEYS[1], 'l')
local now = tonumber(KEYS[2])
local nrcpt = tonumber(KEYS[6])
local leak_rate = tonumber(KEYS[3])
local max_burst = tonumber(KEYS[4])
local prefix = KEYS[1]
local dynr, dynb, leaked = 0, 0, 0
if not last then
-- New bucket
redis.call('HMSET', KEYS[1], 'l', KEYS[2], 'b', '0', 'dr', '10000', 'db', '10000', 'p', tostring(nrcpt))
redis.call('EXPIRE', KEYS[1], KEYS[5])
redis.call('HMSET', prefix, 'l', tostring(now), 'b', '0', 'dr', '10000', 'db', '10000', 'p', tostring(nrcpt))
redis.call('EXPIRE', prefix, KEYS[5])
return {0, '0', '1', '1', '0'}
end

last = tonumber(last)
local burst,pending = unpack(redis.call('HMGET', KEYS[1], 'b', 'p'))

local burst,pending = unpack(redis.call('HMGET', prefix, 'b', 'p'))
burst,pending = tonumber(burst or '0'),tonumber(pending or '0')
-- Sanity to avoid races
if burst < 0 then burst = 0 end
if pending < 0 then pending = 0 end
pending = pending + nrcpt -- this message
-- Perform leak
if burst + pending > 0 then
if burst > 0 and last < tonumber(KEYS[2]) then
local rate = tonumber(KEYS[3])
dynr = tonumber(redis.call('HGET', KEYS[1], 'dr')) / 10000.0
-- If we have any time passed
if burst > 0 and last < now then
dynr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000.0
if dynr == 0 then dynr = 0.0001 end
rate = rate * dynr
leaked = ((now - last) * rate)
leak_rate = leak_rate * dynr
leaked = ((now - last) * leak_rate)
if leaked > burst then leaked = burst end
burst = burst - leaked
redis.call('HINCRBYFLOAT', KEYS[1], 'b', -(leaked))
redis.call('HSET', KEYS[1], 'l', KEYS[2])
redis.call('HINCRBYFLOAT', prefix, 'b', -(leaked))
redis.call('HSET', prefix, 'l', tostring(now))
end

dynb = tonumber(redis.call('HGET', KEYS[1], 'db')) / 10000.0
dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0
if dynb == 0 then dynb = 0.0001 end

burst = burst + pending
if burst > 0 and (burst + tonumber(KEYS[6])) > tonumber(KEYS[4]) * dynb then
if burst > 0 and (burst + nrcpt) > max_burst * dynb then
return {1, tostring(burst - pending), tostring(dynr), tostring(dynb), tostring(leaked)}
end
-- Increase pending if we allow ratelimit
redis.call('HINCRBY', KEYS[1], 'p', nrcpt)
redis.call('HINCRBY', prefix, 'p', nrcpt)
else
burst = 0
redis.call('HMSET', KEYS[1], 'b', '0', 'p', tostring(nrcpt))
redis.call('HMSET', prefix, 'b', '0', 'p', tostring(nrcpt))
end

return {0, tostring(burst), tostring(dynr), tostring(dynb), tostring(leaked)}
52 changes: 27 additions & 25 deletions lualib/redis_scripts/ratelimit_update.lua
View file
Expand Up @@ -13,68 +13,70 @@
-- p - messages pending (must be decreased by 1)
-- dr - current dynamic rate multiplier
-- db - current dynamic burst multiplier

local last = redis.call('HGET', KEYS[1], 'l')
local prefix = KEYS[1]
local last = redis.call('HGET', prefix, 'l')
local now = tonumber(KEYS[2])
local nrcpt = tonumber(KEYS[8])
if not last then
-- New bucket (why??)
redis.call('HMSET', KEYS[1], 'l', KEYS[2], 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0')
redis.call('EXPIRE', KEYS[1], KEYS[7])
redis.call('HMSET', prefix, 'l', tostring(now), 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0')
redis.call('EXPIRE', prefix, KEYS[7])
return {1, 1, 1}
end

local dr, db = 1.0, 1.0

if tonumber(KEYS[5]) > 1 then
local max_dr = tonumber(KEYS[5])

if max_dr > 1 then
local rate_mult = tonumber(KEYS[3])
local rate_limit = tonumber(KEYS[5])
dr = tonumber(redis.call('HGET', KEYS[1], 'dr')) / 10000
dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000

if rate_mult > 1.0 and dr < rate_limit then
if rate_mult > 1.0 and dr < max_dr then
dr = dr * rate_mult
if dr > 0.0001 then
redis.call('HSET', KEYS[1], 'dr', tostring(math.floor(dr * 10000)))
redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
else
redis.call('HSET', KEYS[1], 'dr', '1')
redis.call('HSET', prefix, 'dr', '1')
end
elseif rate_mult < 1.0 and dr > (1.0 / rate_limit) then
elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then
dr = dr * rate_mult
if dr > 0.0001 then
redis.call('HSET', KEYS[1], 'dr', tostring(math.floor(dr * 10000)))
redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000)))
else
redis.call('HSET', KEYS[1], 'dr', '1')
redis.call('HSET', prefix, 'dr', '1')
end
end
end

if tonumber(KEYS[6]) > 1 then
local max_db = tonumber(KEYS[6])
if max_db > 1 then
local rate_mult = tonumber(KEYS[4])
local rate_limit = tonumber(KEYS[6])
db = tonumber(redis.call('HGET', KEYS[1], 'db')) / 10000
db = tonumber(redis.call('HGET', prefix, 'db')) / 10000

if rate_mult > 1.0 and db < rate_limit then
if rate_mult > 1.0 and db < max_db then
db = db * rate_mult
if db > 0.0001 then
redis.call('HSET', KEYS[1], 'db', tostring(math.floor(db * 10000)))
redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
else
redis.call('HSET', KEYS[1], 'db', '1')
redis.call('HSET', prefix, 'db', '1')
end
elseif rate_mult < 1.0 and db > (1.0 / rate_limit) then
elseif rate_mult < 1.0 and db > (1.0 / max_db) then
db = db * rate_mult
if db > 0.0001 then
redis.call('HSET', KEYS[1], 'db', tostring(math.floor(db * 10000)))
redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000)))
else
redis.call('HSET', KEYS[1], 'db', '1')
redis.call('HSET', prefix, 'db', '1')
end
end
end

local burst,pending = unpack(redis.call('HMGET', KEYS[1], 'b', 'p'))
local burst,pending = unpack(redis.call('HMGET', prefix, 'b', 'p'))
burst,pending = tonumber(burst or '0'),tonumber(pending or '0')
if burst < 0 then burst = nrcpt else burst = burst + nrcpt end
if pending < nrcpt then pending = 0 else pending = pending - nrcpt end

redis.call('HMSET', KEYS[1], 'b', tostring(burst), 'p', tostring(pending), 'l', KEYS[2])
redis.call('EXPIRE', KEYS[1], KEYS[7])
redis.call('HMSET', prefix, 'b', tostring(burst), 'p', tostring(pending), 'l', KEYS[2])
redis.call('EXPIRE', prefix, KEYS[7])

return {tostring(burst), tostring(dr), tostring(db)}

0 comments on commit 5d8aa83

Please sign in to comment.