diff --git a/apisix/plugins/limit-count/limit-count-redis-cluster.lua b/apisix/plugins/limit-count/limit-count-redis-cluster.lua index 2be5caec637c..a2da8bd61f26 100644 --- 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 = { name = "apisix-redis-cluster", @@ -77,56 +85,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 index 24621437dffd..7958bdc840d9 100644 --- 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