[Project] Enable compatibility with the existing buckets

2 files changed, 56 insertions, 40 deletions

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