From c4386c6ab0d649d3b999f0cf98f6fccea9212ce5 Mon Sep 17 00:00:00 2001
From: LeftTry <lerest.go@gmail.com>
Date: Sun, 29 Sep 2024 12:23:46 +0600
Subject: [PATCH] [Feature] Add LRU cache for last filled ratelimit buckets

---
 lualib/redis_scripts/ratelimit_check.lua | 5 +++++
 src/plugins/lua/ratelimit.lua            | 5 ++++-
 2 files changed, 9 insertions(+), 1 deletion(-)

diff --git a/lualib/redis_scripts/ratelimit_check.lua b/lualib/redis_scripts/ratelimit_check.lua
index f24e0daf0..f066610a5 100644
--- a/lualib/redis_scripts/ratelimit_check.lua
+++ b/lualib/redis_scripts/ratelimit_check.lua
@@ -31,6 +31,8 @@ local leak_rate = tonumber(KEYS[3])
 local max_burst = tonumber(KEYS[4])
 local prefix = KEYS[1]
 local enable_dynamic = KEYS[7] == 'true'
+local cache_prefix = KEYS[8]
+local max_cache_size = KEYS[9]
 local dynr, dynb, leaked = 0, 0, 0
 if not last then
   -- New bucket
@@ -83,6 +85,9 @@ if burst + pending > 0 then
 
   burst = burst + pending
   if burst > 0 and burst > max_burst * dynb then
+    redis.call('ZREMRANGEBYRANK', cache_prefix, 0, -(max_cache_size + 1))
+    redis.call('ZINCRBY', cache_prefix, 1, prefix)
+
     return { 1, tostring(burst - pending), tostring(dynr), tostring(dynb), tostring(leaked) }
   end
   -- Increase pending if we allow ratelimit
diff --git a/src/plugins/lua/ratelimit.lua b/src/plugins/lua/ratelimit.lua
index 168d8d63a..c0783eac8 100644
--- a/src/plugins/lua/ratelimit.lua
+++ b/src/plugins/lua/ratelimit.lua
@@ -41,6 +41,8 @@ local settings = {
   -- Do not check ratelimits for these recipients
   whitelisted_rcpts = { 'postmaster', 'mailer-daemon' },
   prefix = 'RL',
+  cache_prefix = 'RL_cache_prefix',
+  max_cache_size = 30,
   -- If enabled, we apply dynamic rate limiting based on the verdict
   dynamic_rate_limit = false,
   ham_factor_rate = 1.01,
@@ -455,7 +457,8 @@ local function ratelimit_cb(task)
           { 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(dyn_rate_enabled) })
+            tostring(settings.expire), tostring(bincr), tostring(dyn_rate_enabled),
+            settings.cache_prefix, settings.max_cache_size })
     end
   end
 end
-- 
2.39.5