From 5d8aa83eb61e59a2360e7d70e971217f8567d13a Mon Sep 17 00:00:00 2001 From: Vsevolod Stakhov Date: Sat, 25 Mar 2023 13:19:35 +0000 Subject: [Minor] Improve style in the scripts --- lualib/redis_scripts/ratelimit_check.lua | 33 +++++++++++--------- lualib/redis_scripts/ratelimit_update.lua | 52 ++++++++++++++++--------------- 2 files changed, 45 insertions(+), 40 deletions(-) (limited to 'lualib/redis_scripts') diff --git a/lualib/redis_scripts/ratelimit_check.lua b/lualib/redis_scripts/ratelimit_check.lua index 2b2af11bf..e43d34ce1 100644 --- a/lualib/redis_scripts/ratelimit_check.lua +++ b/lualib/redis_scripts/ratelimit_check.lua @@ -16,16 +16,19 @@ 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 @@ -33,30 +36,30 @@ 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)} \ No newline at end of file diff --git a/lualib/redis_scripts/ratelimit_update.lua b/lualib/redis_scripts/ratelimit_update.lua index 682ddd0c6..01ab4e051 100644 --- a/lualib/redis_scripts/ratelimit_update.lua +++ b/lualib/redis_scripts/ratelimit_update.lua @@ -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)} \ No newline at end of file -- cgit v1.2.3