test/apisix
1
0
Fork 0
mirror of https://gitee.com/iresty/apisix.git synced 2024-12-14 08:51:29 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: yuchenghao <yuchenghao@didachuxing.com>
This commit is contained in:
ChenghaoYu 2021-04-08 10:16:30 +08:00 committed by GitHub
parent 7dc60f1068
commit 252af4143d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 22 additions and 91 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

53
apisix/plugins/limit-count/limit-count-redis-cluster.lua
View File

@ -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
local remaining, err = red:eval(script, 1, key, limit, window)
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
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

60
apisix/plugins/limit-count/limit-count-redis.lua
View File

@ -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
remaining, err = red:eval(script, 1, key, limit, window)
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
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