Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 5 commits into from
Apr 8, 2021
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 11 additions & 46 deletions apisix/plugins/limit-count/limit-count-redis-cluster.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -72,62 +71,28 @@ function _M.new(plugin_name, limit, window, conf)
return setmetatable(self, mt)
end

local script = "if redis.call('ttl',KEYS[1]) < 0 then "
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be better to use exists instead of ttl?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think ttl is better. ttl will set the expire time on the key which is permanent. Both are O(1) Time complexity.

.. "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.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
local remaining,err = red:eval(script,1,key,limit,window)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Need a space around the operator. Please fix other similar places.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok. thanks


return 0, limit -1
end

local ok, err = lock:unlock()
if not ok then
return false, "failed to unlock: " .. err
end
if err then
return nil,err
end

remaining, err = red:incrby(key, -1)
if not remaining then
return nil, err
if remaining <0 then
return nil,"rejected"
end

if remaining < 0 then
return nil, "rejected"
end

return 0, remaining
return 0,remaining
end


Expand Down
68 changes: 13 additions & 55 deletions apisix/plugins/limit-count/limit-count-redis.lua
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -30,6 +29,13 @@ 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)

Expand Down Expand Up @@ -76,65 +82,17 @@ function _M.incoming(self, key)

local limit = self.limit
local window = self.window
local remaining
key = self.plugin_name .. tostring(key)
local remaining, err = red:eval(script,1,key,limit,window)

-- todo: test case
local ret, err = red:ttl(key)
if not ret then
return false, "failed to get redis `" .. key .."` ttl: " .. err
if err then
return nil,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

local ok, err = red:set_keepalive(10000, 100)
if not ok then
return nil, err
if remaining <0 then
return nil,"rejected"
end

if remaining < 0 then
return nil, "rejected"
end

return 0, remaining
return 0,remaining
end


Expand Down