diff --git a/apisix/plugins/limit-req/util.lua b/apisix/plugins/limit-req/util.lua index 282c04cf9059..fb63e242c7e1 100644 --- a/apisix/plugins/limit-req/util.lua +++ b/apisix/plugins/limit-req/util.lua @@ -14,60 +14,66 @@ -- See the License for the specific language governing permissions and -- limitations under the License. -- -local math = require "math" -local abs = math.abs -local max = math.max local ngx_now = ngx.now -local ngx_null = ngx.null local tonumber = tonumber +local core = require("apisix.core") local _M = {version = 0.1} +local script = core.string.compress_script([=[ + local state_key = KEYS[1] -- state_key (hash), fields: "excess", "last" + local rate = tonumber(ARGV[1]) -- req/s + local now = tonumber(ARGV[2]) -- ms + local burst = tonumber(ARGV[3]) -- req/s + local commit = tonumber(ARGV[4]) -- 1/0 + + local vals = redis.call("HMGET", state_key, "excess", "last") + local prev_excess = tonumber(vals[1] or "0") + local prev_last = tonumber(vals[2] or "0") + + local new_excess + if prev_last > 0 then + local elapsed = math.abs(now - prev_last) + new_excess = math.max(prev_excess - rate * (elapsed) / 1000 + 1000, 0) + else + new_excess = 0 + end + + if new_excess > burst then + return {0, new_excess} + end + + if commit == 1 then + redis.call("HMSET", state_key, "excess", new_excess, "last", now) + redis.call("EXPIRE", state_key, 60) + end + + return {1, new_excess} +]=]) + + -- the "commit" argument controls whether should we record the event in shm. function _M.incoming(self, red, key, commit) local rate = self.rate local now = ngx_now() * 1000 - key = "limit_req" .. ":" .. key - local excess_key = key .. "excess" - local last_key = key .. "last" + local state_key = "limit_req:{" .. key .. "}:state" - local excess, err = red:get(excess_key) - if err then - return nil, err - end - local last, err = red:get(last_key) - if err then + local commit_flag = commit and "1" or "0" + + local res, err = red:eval(script, 1, state_key, + rate, now, self.burst, commit_flag) + if not res then return nil, err end - if excess ~= ngx_null and last ~= ngx_null then - excess = tonumber(excess) - last = tonumber(last) - local elapsed = now - last - excess = max(excess - rate * abs(elapsed) / 1000 + 1000, 0) - - if excess > self.burst then - return nil, "rejected" - end - else - excess = 0 - end + local allowed = tonumber(res[1]) == 1 + local excess = tonumber(res[2]) or 0 - if commit then - local ok - local err - ok, err = red:set(excess_key, excess) - if not ok then - return nil, err - end - - ok, err = red:set(last_key, now) - if not ok then - return nil, err - end + if not allowed then + return nil, "rejected" end -- return the delay in seconds, as well as excess diff --git a/t/plugin/limit-req-redis-cluster.t b/t/plugin/limit-req-redis-cluster.t index 4c36c2200294..05db77a76c3f 100644 --- a/t/plugin/limit-req-redis-cluster.t +++ b/t/plugin/limit-req-redis-cluster.t @@ -603,3 +603,121 @@ passed GET /t --- response_body eval qr/property \"rate\" validation failed: expected 0 to be greater than 0/ + + + +=== TEST 22: verify atomic Redis cluster operations with hash key structure +--- config + location /t { + content_by_lua_block { + local redis_cluster = require("apisix.utils.rediscluster") + local conf = { + redis_cluster_name = "test", + redis_cluster_nodes = { + "127.0.0.1:5000", + "127.0.0.1:5002" + } + } + local red_c, err = redis_cluster.new(conf, "plugin-limit-req-redis-cluster-slot-lock") + if not red_c then + ngx.say("Failed to create Redis cluster client: ", err) + return + end + + -- Clean up any existing keys + red_c:del("limit_req:{test_key}:state") + + -- Test the new hash-based key structure + local util = require("apisix.plugins.limit-req.util") + local limiter = { + rate = 10, -- 10 req/s + burst = 1000 -- 1000 req/s burst + } + + -- First request should succeed (use non-pipeline client) + local delay, excess = util.incoming(limiter, red_c, "test_key", true) + if delay then + ngx.say("first request: delay=", delay, " excess=", excess) + else + ngx.say("first request failed: ", excess) + end + + -- Verify the Redis hash was created with correct key format + local vals, err = red_c:hmget("limit_req:{test_key}:state", "excess", "last") + if vals and vals[1] and vals[2] then + ngx.say("hash key created: excess=", vals[1], " last=", vals[2]) + else + ngx.say("hash key not found") + end + + -- Verify TTL was set + local ttl = red_c:ttl("limit_req:{test_key}:state") + if ttl and ttl > 0 then + ngx.say("TTL set: ", ttl, " seconds") + else + ngx.say("TTL not set") + end + + -- Clean up + red_c:del("limit_req:{test_key}:state") + } + } +--- request +GET /t +--- response_body_like +first request: delay=\d+\.?\d* excess=\d+\.?\d* +hash key created: excess=\d+ last=\d+ +TTL set: \d+ seconds + + + +=== TEST 23: verify atomic behavior prevents race conditions in cluster +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "plugins": { + "limit-req": { + "rate": 1, + "burst": 0, + "rejected_code": 503, + "key": "remote_addr", + "policy": "redis-cluster", + "redis_cluster_name": "test", + "redis_cluster_nodes": [ + "127.0.0.1:5000", + "127.0.0.1:5002" + ] + } + }, + "upstream": { + "nodes": { + "127.0.0.1:1980": 1 + }, + "type": "roundrobin" + }, + "uri": "/hello" + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- request +GET /t +--- response_body +passed + + + +=== TEST 24: test atomic rate limiting with rapid requests in cluster +--- pipelined_requests eval +["GET /hello", "GET /hello"] +--- error_code eval +[200, 503] diff --git a/t/plugin/limit-req-redis.t b/t/plugin/limit-req-redis.t index 84664b7a2db2..9b3493f6de49 100644 --- a/t/plugin/limit-req-redis.t +++ b/t/plugin/limit-req-redis.t @@ -651,3 +651,114 @@ passed GET /t --- response_body eval qr/property \"rate\" validation failed: expected 0 to be greater than 0/ + + + +=== TEST 24: verify atomic Redis operations with hash key structure +--- config + location /t { + content_by_lua_block { + local redis = require "resty.redis" + local red = redis:new() + red:set_timeout(1000) + + local ok, err = red:connect("127.0.0.1", 6379) + if not ok then + ngx.say("failed to connect: ", err) + return + end + + -- Clean up any existing keys + red:del("limit_req:{test_key}:state") + + -- Test the new hash-based key structure + local util = require("apisix.plugins.limit-req.util") + local limiter = { + rate = 10, -- 10 req/s + burst = 1000 -- 1000 req/s burst + } + + -- First request should succeed + local delay, excess = util.incoming(limiter, red, "test_key", true) + if delay then + ngx.say("first request: delay=", delay, " excess=", excess) + else + ngx.say("first request failed: ", excess) + end + + -- Verify the Redis hash was created with correct key format + local vals = red:hmget("limit_req:{test_key}:state", "excess", "last") + if vals[1] and vals[2] then + ngx.say("hash key created: excess=", vals[1], " last=", vals[2]) + else + ngx.say("hash key not found") + end + + -- Verify TTL was set + local ttl = red:ttl("limit_req:{test_key}:state") + if ttl and ttl > 0 then + ngx.say("TTL set: ", ttl, " seconds") + else + ngx.say("TTL not set") + end + + -- Clean up + red:del("limit_req:{test_key}:state") + red:close() + } + } +--- request +GET /t +--- response_body_like +first request: delay=\d+\.?\d* excess=\d+\.?\d* +hash key created: excess=\d+ last=\d+ +TTL set: \d+ seconds + + + +=== TEST 25: verify atomic behavior prevents race conditions +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "plugins": { + "limit-req": { + "rate": 1, + "burst": 0, + "rejected_code": 503, + "key": "remote_addr", + "policy": "redis", + "redis_host": "127.0.0.1" + } + }, + "upstream": { + "nodes": { + "127.0.0.1:1980": 1 + }, + "type": "roundrobin" + }, + "uri": "/hello" + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- request +GET /t +--- response_body +passed + + + +=== TEST 26: test atomic rate limiting with rapid requests +--- pipelined_requests eval +["GET /hello", "GET /hello"] +--- error_code eval +[200, 503]