From: Vsevolod Stakhov Date: Tue, 9 Jul 2024 13:22:31 +0000 (+0100) Subject: [Project] Enable compatibility with the existing buckets X-Git-Tag: 3.9.0~4^2~3 X-Git-Url: https://source.dussan.org/?a=commitdiff_plain;h=87702745e128977d238284b2a69415bcf94974e6;p=rspamd.git [Project] Enable compatibility with the existing buckets --- diff --git a/lualib/redis_scripts/ratelimit_check.lua b/lualib/redis_scripts/ratelimit_check.lua index d39cdf148..f24e0daf0 100644 --- a/lualib/redis_scripts/ratelimit_check.lua +++ b/lualib/redis_scripts/ratelimit_check.lua @@ -7,6 +7,7 @@ -- KEYS[4]: The maximum allowed burst -- KEYS[5]: The expiration time for a bucket -- KEYS[6]: The number of recipients for the message +-- KEYS[7]: Enable dynamic ratelimits -- Redis keys used: -- l: Last hit (time in milliseconds) @@ -29,6 +30,7 @@ local nrcpt = tonumber(KEYS[6]) local leak_rate = tonumber(KEYS[3]) local max_burst = tonumber(KEYS[4]) local prefix = KEYS[1] +local enable_dynamic = KEYS[7] == 'true' local dynr, dynb, leaked = 0, 0, 0 if not last then -- New bucket @@ -52,9 +54,13 @@ pending = pending + nrcpt -- this message if burst + pending > 0 then -- 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 + if enable_dynamic then + dynr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000.0 + if dynr == 0 then + dynr = 0.0001 + end + else + dynr = 1.0 end leak_rate = leak_rate * dynr leaked = ((now - last) * leak_rate) @@ -66,9 +72,13 @@ if burst + pending > 0 then redis.call('HSET', prefix, 'l', tostring(now)) end - dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0 - if dynb == 0 then - dynb = 0.0001 + if enable_dynamic then + dynb = tonumber(redis.call('HGET', prefix, 'db')) / 10000.0 + if dynb == 0 then + dynb = 0.0001 + end + else + dynb = 1.0 end burst = burst + pending diff --git a/lualib/redis_scripts/ratelimit_update.lua b/lualib/redis_scripts/ratelimit_update.lua index caee8fb31..8b7a934dc 100644 --- a/lualib/redis_scripts/ratelimit_update.lua +++ b/lualib/redis_scripts/ratelimit_update.lua @@ -9,12 +9,14 @@ -- KEYS[6] - max_burst_rate: The maximum allowed value for the dynamic burst multiplier. -- KEYS[7] - expire: The expiration time for the Redis key storing the bucket information, in seconds. -- KEYS[8] - number_of_recipients: The number of requests to be allowed (or the increase rate). +-- KEYS[9] - Enable dynamic ratelimits -- 1. Retrieve the last hit time and initialize variables local prefix = KEYS[1] local last = redis.call('HGET', prefix, 'l') local now = tonumber(KEYS[2]) local nrcpt = tonumber(KEYS[8]) +local enable_dynamic = KEYS[9] == 'true' if not last then -- 2. Initialize a new bucket if the last hit time is not found (must not happen) redis.call('HMSET', prefix, 'l', tostring(now), 'b', tostring(nrcpt), 'dr', '10000', 'db', '10000', 'p', '0') @@ -25,48 +27,52 @@ end -- 3. Update the dynamic rate multiplier based on input parameters local dr, db = 1.0, 1.0 -local max_dr = tonumber(KEYS[5]) +if enable_dynamic then + local max_dr = tonumber(KEYS[5]) -if max_dr > 1 then - local rate_mult = tonumber(KEYS[3]) - dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000 + if max_dr > 1 then + local rate_mult = tonumber(KEYS[3]) + dr = tonumber(redis.call('HGET', prefix, 'dr')) / 10000 - if rate_mult > 1.0 and dr < max_dr then - dr = dr * rate_mult - if dr > 0.0001 then - redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) - else - redis.call('HSET', prefix, 'dr', '1') - end - elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then - dr = dr * rate_mult - if dr > 0.0001 then - redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) - else - redis.call('HSET', prefix, 'dr', '1') + if rate_mult > 1.0 and dr < max_dr then + dr = dr * rate_mult + if dr > 0.0001 then + redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) + else + redis.call('HSET', prefix, 'dr', '1') + end + elseif rate_mult < 1.0 and dr > (1.0 / max_dr) then + dr = dr * rate_mult + if dr > 0.0001 then + redis.call('HSET', prefix, 'dr', tostring(math.floor(dr * 10000))) + else + redis.call('HSET', prefix, 'dr', '1') + end end end end -- 4. Update the dynamic burst multiplier based on input parameters -local max_db = tonumber(KEYS[6]) -if max_db > 1 then - local rate_mult = tonumber(KEYS[4]) - db = tonumber(redis.call('HGET', prefix, 'db')) / 10000 +if enable_dynamic then + local max_db = tonumber(KEYS[6]) + if max_db > 1 then + local rate_mult = tonumber(KEYS[4]) + db = tonumber(redis.call('HGET', prefix, 'db')) / 10000 - if rate_mult > 1.0 and db < max_db then - db = db * rate_mult - if db > 0.0001 then - redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) - else - redis.call('HSET', prefix, 'db', '1') - end - elseif rate_mult < 1.0 and db > (1.0 / max_db) then - db = db * rate_mult - if db > 0.0001 then - redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) - else - redis.call('HSET', prefix, 'db', '1') + if rate_mult > 1.0 and db < max_db then + db = db * rate_mult + if db > 0.0001 then + redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) + else + redis.call('HSET', prefix, 'db', '1') + end + elseif rate_mult < 1.0 and db > (1.0 / max_db) then + db = db * rate_mult + if db > 0.0001 then + redis.call('HSET', prefix, 'db', tostring(math.floor(db * 10000))) + else + redis.call('HSET', prefix, 'db', '1') + end end end end diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua index 470ac5f07..c3602f915 100644 --- a/src/plugins/lua/ratelimit.lua +++ b/src/plugins/lua/ratelimit.lua @@ -365,17 +365,33 @@ local function make_prefix(redis_key, name, bucket) -- Fill defaults -- If settings.dynamic_rate_limit is false, then the default dynamic rate limits are 1.0 -- We always allow per-bucket overrides of the dyn rate limits + + local seen_specific_dyn_rate = false + if not bucket.spam_factor_rate then bucket.spam_factor_rate = settings.dynamic_rate_limit and settings.spam_factor_rate or 1.0 + else + seen_specific_dyn_rate = true end if not bucket.ham_factor_rate then bucket.ham_factor_rate = settings.dynamic_rate_limit and settings.ham_factor_rate or 1.0 + else + seen_specific_dyn_rate = true end if not bucket.spam_factor_burst then bucket.spam_factor_burst = settings.dynamic_rate_limit and settings.spam_factor_burst or 1.0 + else + seen_specific_dyn_rate = true end if not bucket.ham_factor_burst then bucket.ham_factor_burst = settings.dynamic_rate_limit and settings.ham_factor_burst or 1.0 + else + seen_specific_dyn_rate = true + end + + if seen_specific_dyn_rate then + -- Use if afterwards in case we don't use global dyn rates + bucket.specific_dyn_rate = true end return { @@ -555,13 +571,15 @@ local function ratelimit_cb(task) bincr = 1 end + local dyn_rate_enabled = settings.dynamic_rate_limit or bucket.specific_dyn_rate + lua_util.debugm(N, task, "check limit %s:%s -> %s (%s/%s)", value.name, pr, value.hash, bucket.burst, bucket.rate) lua_redis.exec_redis_script(bucket_check_id, { key = value.hash, task = task, is_write = true }, gen_check_cb(pr, bucket, value.name, value.hash), { value.hash, tostring(now), tostring(rate), tostring(bucket.burst), - tostring(settings.expire), tostring(bincr) }) + tostring(settings.expire), tostring(bincr), tostring(dyn_rate_enabled) }) end end end @@ -661,12 +679,14 @@ local function ratelimit_update_cb(task) bincr = 1 end + local dyn_rate_enabled = settings.dynamic_rate_limit or bucket.specific_dyn_rate + lua_redis.exec_redis_script(bucket_update_id, { key = v.hash, task = task, is_write = true }, update_bucket_cb, { v.hash, tostring(now), tostring(mult_rate), tostring(mult_burst), tostring(settings.max_rate_mult), tostring(settings.max_bucket_mult), - tostring(settings.expire), tostring(bincr) }) + tostring(settings.expire), tostring(bincr), tostring(dyn_rate_enabled) }) end end end