feat: add limit-count sliding window for redis with tests

Signed-off-by: Sihyeon Jang <[email protected]>

Author: sihyeonn

apisix/plugins/limit-count/limit-count-redis.lua

@@ -21,15 +21,15 @@ local setmetatable = setmetatable
 local tostring = tostring
 
 
-local _M = {version = 0.3}
+local _M = {version = 0.4}
 
 
 local mt = {
     __index = _M
 }
 
 
-local script = core.string.compress_script([=[
+local script_fixed = core.string.compress_script([=[
     assert(tonumber(ARGV[3]) >= 1, "cost must be at least 1")
     local ttl = redis.call('ttl', KEYS[1])
     if ttl < 0 then
@@ -40,12 +40,61 @@ local script = core.string.compress_script([=[
 ]=])
 
 
-function _M.new(plugin_name, limit, window, conf)
+local script_sliding = core.string.compress_script([=[
+    assert(tonumber(ARGV[3]) >= 1, "cost must be at least 1")
+
+    local now = tonumber(ARGV[1])
+    local window = tonumber(ARGV[2])
+    local limit = tonumber(ARGV[3])
+    local cost = tonumber(ARGV[4])
+
+    local window_start = now - window
+
+    -- remove events outside of the window
+    redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, window_start)
+
+    local current = redis.call('ZCARD', KEYS[1])
+
+    if current + cost > limit then
+        local earliest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
+        local reset = 0
+        if #earliest == 2 then
+            reset = earliest[2] + window - now
+            if reset < 0 then
+                reset = 0
+            end
+        end
+        return {-1, reset}
+    end
+
+    for i = 1, cost do
+        redis.call('ZADD', KEYS[1], now, now .. ':' .. i)
+    end
+
+    redis.call('PEXPIRE', KEYS[1], window)
+
+    local remaining = limit - (current + cost)
+
+    local earliest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
+    local reset = 0
+    if #earliest == 2 then
+        reset = earliest[2] + window - now
+        if reset < 0 then
+            reset = 0
+        end
+    end
+
+    return {remaining, reset}
+]=])
+
+
+function _M.new(plugin_name, limit, window, window_type, conf)
     assert(limit > 0 and window > 0)
 
     local self = {
         limit = limit,
         window = window,
+        window_type = window_type or "fixed",
         conf = conf,
         plugin_name = plugin_name,
     }
@@ -59,13 +108,22 @@ function _M.incoming(self, key, cost)
         return red, err, 0
     end
 
-    local limit = self.limit
-    local window = self.window
-    local res
     key = self.plugin_name .. tostring(key)
 
     local ttl = 0
-    res, err = red:eval(script, 1, key, limit, window, cost or 1)
+    local limit = self.limit
+    local c = cost or 1
+    local res
+
+    if self.window_type == "sliding" then
+        local now = ngx.now() * 1000
+        local window = self.window * 1000
+
+        res, err = red:eval(script_sliding, 1, key, now, window, limit, c)
+    else
+        local window = self.window
+        res, err = red:eval(script_fixed, 1, key, limit, window, c)
+    end
 
     if err then
         return nil, err, ttl
@@ -74,14 +132,15 @@ function _M.incoming(self, key, cost)
     local remaining = res[1]
     ttl = res[2]
 
-    local ok, err = red:set_keepalive(10000, 100)
+    local ok, err2 = red:set_keepalive(10000, 100)
     if not ok then
-        return nil, err, ttl
+        return nil, err2, ttl
     end
 
     if remaining < 0 then
         return nil, "rejected", ttl
     end
+
     return 0, remaining, ttl
 end