fix: ensure automic operation in limit-count plugin (#3991)

Co-authored-by: yuchenghao <[email protected]>
apache · Apr 8, 2021 · 252af41 · 252af41
1 parent 7dc60f1
commit 252af41
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 91 deletions.
diff --git a/apisix/plugins/limit-count/limit-count-redis-cluster.lua b/apisix/plugins/limit-count/limit-count-redis-cluster.lua
@@ -17,7 +17,6 @@
 
 local rediscluster = require("resty.rediscluster")
 local core = require("apisix.core")
-local resty_lock = require("resty.lock")
 local setmetatable = setmetatable
 local tostring = tostring
 local ipairs = ipairs
@@ -30,6 +29,15 @@ local mt = {
 }
 
 
+local script = [=[
+    if redis.call('ttl', KEYS[1]) < 0 then
+        redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
+        return ARGV[1] - 1
+    end
+    return redis.call('incrby', KEYS[1], -1)
+]=]
+
+
 local function new_redis_cluster(conf)
     local config = {
         -- can set different name for different redis cluster
@@ -78,56 +86,17 @@ function _M.incoming(self, key)
     local red = self.red_cli
     local limit = self.limit
     local window = self.window
-    local remaining
     key = self.plugin_name .. tostring(key)
 
-    local ret, err = red:ttl(key)
-    if not ret then
-        return false, "failed to get redis `" .. key .."` ttl: " .. err
-    end
-
-    core.log.info("ttl key: ", key, " ret: ", ret, " err: ", err)
-    if ret < 0 then
-        local lock, err = resty_lock:new("plugin-limit-count")
-        if not lock then
-            return false, "failed to create lock: " .. err
-        end
-
-        local elapsed, err = lock:lock(key)
-        if not elapsed then
-            return false, "failed to acquire the lock: " .. err
-        end
-
-        ret = red:ttl(key)
-        if ret < 0 then
-            local ok, err = lock:unlock()
-            if not ok then
-                return false, "failed to unlock: " .. err
-            end
-
-            ret, err = red:set(key, limit -1, "EX", window)
-            if not ret then
-                return nil, err
-            end
-
-            return 0, limit -1
-        end
-
-        local ok, err = lock:unlock()
-        if not ok then
-            return false, "failed to unlock: " .. err
-        end
-    end
+    local remaining, err = red:eval(script, 1, key, limit, window)
 
-    remaining, err = red:incrby(key, -1)
-    if not remaining then
+    if err then
         return nil, err
     end
 
     if remaining < 0 then
         return nil, "rejected"
     end
-
     return 0, remaining
 end
 

diff --git a/apisix/plugins/limit-count/limit-count-redis.lua b/apisix/plugins/limit-count/limit-count-redis.lua
@@ -16,7 +16,6 @@
 --
 local redis_new = require("resty.redis").new
 local core = require("apisix.core")
-local resty_lock = require("resty.lock")
 local assert = assert
 local setmetatable = setmetatable
 local tostring = tostring
@@ -30,6 +29,15 @@ local mt = {
 }
 
 
+local script = [=[
+    if redis.call('ttl', KEYS[1]) < 0 then
+        redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
+        return ARGV[1] - 1
+    end
+    return redis.call('incrby', KEYS[1], -1)
+]=]
+
+
 function _M.new(plugin_name, limit, window, conf)
     assert(limit > 0 and window > 0)
 
@@ -79,61 +87,15 @@ function _M.incoming(self, key)
     local remaining
     key = self.plugin_name .. tostring(key)
 
-    -- todo: test case
-    local ret, err = red:ttl(key)
-    if not ret then
-        return false, "failed to get redis `" .. key .."` ttl: " .. err
-    end
-
-    core.log.info("ttl key: ", key, " ret: ", ret, " err: ", err)
-    if ret < 0 then
-        -- todo: test case
-        local lock, err = resty_lock:new("plugin-limit-count")
-        if not lock then
-            return false, "failed to create lock: " .. err
-        end
-
-        local elapsed, err = lock:lock(key)
-        if not elapsed then
-            return false, "failed to acquire the lock: " .. err
-        end
-
-        ret = red:ttl(key)
-        if ret < 0 then
-            ok, err = lock:unlock()
-            if not ok then
-                return false, "failed to unlock: " .. err
-            end
-
-            limit = limit -1
-            ret, err = red:set(key, limit, "EX", window)
-            if not ret then
-                return nil, err
-            end
-
-            return 0, limit
-        end
-
-        ok, err = lock:unlock()
-        if not ok then
-            return false, "failed to unlock: " .. err
-        end
-    end
-
-    remaining, err = red:incrby(key, -1)
-    if not remaining then
-        return nil, err
-    end
+    remaining, err = red:eval(script, 1, key, limit, window)
 
-    local ok, err = red:set_keepalive(10000, 100)
-    if not ok then
+    if err then
         return nil, err
     end
 
     if remaining < 0 then
         return nil, "rejected"
     end
-
     return 0, remaining
 end