feat: truncate suspicious_ips when it exceed the max_size

lablup · Apr 30, 2024 · be22862 · be22862
1 parent 1773971
commit be22862
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/ai/backend/manager/api/ratelimit.py b/src/ai/backend/manager/api/ratelimit.py
@@ -55,6 +55,12 @@
     -- Add IP to suspicious_ips only if count is greater than score_threshold
     if rolling_count >= score_threshold then
         redis.call('ZADD', 'suspicious_ips', rolling_count, id_value)
+
+        local max_size = 1000
+        local current_size = redis.call('ZCARD', 'suspicious_ips')
+        if current_size > max_size then
+            redis.call('ZREMRANGEBYRANK', 'suspicious_ips', 0, 0)
+        end
     end
 end