Skip to content

Commit 03e4161

Browse files
sihyeonnsihyeonn
committed
feat: add limit-count sliding window for redis-cluster with tests
Signed-off-by: Sihyeon Jang <[email protected]>
1 parent 9da2fcb commit 03e4161

File tree

2 files changed

+253
-6
lines changed

2 files changed

+253
-6
lines changed

‎apisix/plugins/limit-count/limit-count-redis-cluster.lua‎

Lines changed: 65 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,15 +20,15 @@ local core = require("apisix.core")
2020
local setmetatable = setmetatable
2121
local tostring = tostring
2222

23-
local _M = {}
23+
local _M = {version = 0.2}
2424

2525

2626
local mt = {
2727
__index = _M
2828
}
2929

3030

31-
local script = core.string.compress_script([=[
31+
local script_fixed = core.string.compress_script([=[
3232
assert(tonumber(ARGV[3]) >= 1, "cost must be at least 1")
3333
local ttl = redis.call('ttl', KEYS[1])
3434
if ttl < 0 then
@@ -39,7 +39,54 @@ local script = core.string.compress_script([=[
3939
]=])
4040

4141

42-
function _M.new(plugin_name, limit, window, conf)
42+
local script_sliding = core.string.compress_script([=[
43+
assert(tonumber(ARGV[3]) >= 1, "cost must be at least 1")
44+
45+
local now = tonumber(ARGV[1])
46+
local window = tonumber(ARGV[2])
47+
local limit = tonumber(ARGV[3])
48+
local cost = tonumber(ARGV[4])
49+
50+
local window_start = now - window
51+
52+
redis.call('ZREMRANGEBYSCORE', KEYS[1], 0, window_start)
53+
54+
local current = redis.call('ZCARD', KEYS[1])
55+
56+
if current + cost > limit then
57+
local earliest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
58+
local reset = 0
59+
if #earliest == 2 then
60+
reset = earliest[2] + window - now
61+
if reset < 0 then
62+
reset = 0
63+
end
64+
end
65+
return {-1, reset}
66+
end
67+
68+
for i = 1, cost do
69+
redis.call('ZADD', KEYS[1], now, now .. ':' .. i)
70+
end
71+
72+
redis.call('PEXPIRE', KEYS[1], window)
73+
74+
local remaining = limit - (current + cost)
75+
76+
local earliest = redis.call('ZRANGE', KEYS[1], 0, 0, 'WITHSCORES')
77+
local reset = 0
78+
if #earliest == 2 then
79+
reset = earliest[2] + window - now
80+
if reset < 0 then
81+
reset = 0
82+
end
83+
end
84+
85+
return {remaining, reset}
86+
]=])
87+
88+
89+
function _M.new(plugin_name, limit, window, window_type, conf)
4390
local red_cli, err = redis_cluster.new(conf, "plugin-limit-count-redis-cluster-slot-lock")
4491
if not red_cli then
4592
return nil, err
@@ -48,6 +95,7 @@ function _M.new(plugin_name, limit, window, conf)
4895
local self = {
4996
limit = limit,
5097
window = window,
98+
window_type = window_type or "fixed",
5199
conf = conf,
52100
plugin_name = plugin_name,
53101
red_cli = red_cli,
@@ -59,12 +107,23 @@ end
59107

60108
function _M.incoming(self, key, cost)
61109
local red = self.red_cli
62-
local limit = self.limit
63-
local window = self.window
64110
key = self.plugin_name .. tostring(key)
65111

66112
local ttl = 0
67-
local res, err = red:eval(script, 1, key, limit, window, cost or 1)
113+
local limit = self.limit
114+
local c = cost or 1
115+
local res
116+
117+
if self.window_type == "sliding" then
118+
local now = ngx.now() * 1000
119+
local window = self.window * 1000
120+
121+
res, err = red:eval(script_sliding, 1, key, now, window, limit, c)
122+
else
123+
local window = self.window
124+
125+
res, err = red:eval(script_fixed, 1, key, limit, window, c)
126+
end
68127

69128
if err then
70129
return nil, err, ttl

‎t/plugin/limit-count-redis-cluster-sliding.t‎

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one or more
3+
# contributor license agreements. See the NOTICE file distributed with
4+
# this work for additional information regarding copyright ownership.
5+
# The ASF licenses this file to You under the Apache License, Version 2.0
6+
# (the "License"); you may not use this file except in compliance with
7+
# the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
18+
use t::APISIX 'no_plan';
19+
20+
repeat_each(1);
21+
no_long_string();
22+
no_shuffle();
23+
no_root_location();
24+
25+
add_block_preprocessor(sub {
26+
my ($block) = @_;
27+
28+
if (!$block->request) {
29+
$block->set_value("request", "GET /t");
30+
}
31+
32+
if (!$block->error_log && !$block->no_error_log) {
33+
$block->set_value("no_error_log", "[error]\n[alert]");
34+
}
35+
});
36+
37+
run_tests;
38+
39+
__DATA__
40+
41+
=== TEST 1: redis-cluster policy with sliding window - basic N per window
42+
--- config
43+
location /t {
44+
content_by_lua_block {
45+
local t = require("lib.test_admin").test
46+
local code, body = t('/apisix/admin/routes/1',
47+
ngx.HTTP_PUT,
48+
[[{
49+
"uri": "/hello",
50+
"plugins": {
51+
"limit-count": {
52+
"count": 2,
53+
"time_window": 2,
54+
"window_type": "sliding",
55+
"rejected_code": 503,
56+
"key": "remote_addr",
57+
"policy": "redis-cluster",
58+
"redis_cluster_nodes": [
59+
"127.0.0.1:5000",
60+
"127.0.0.1:5001"
61+
],
62+
"redis_cluster_name": "redis-cluster-1"
63+
}
64+
},
65+
"upstream": {
66+
"nodes": {
67+
"127.0.0.1:1980": 1
68+
},
69+
"type": "roundrobin"
70+
}
71+
}]]
72+
)
73+
74+
if code >= 300 then
75+
ngx.status = code
76+
end
77+
ngx.say(body)
78+
}
79+
}
80+
--- response_body
81+
passed
82+
83+
84+
=== TEST 2: redis-cluster policy with sliding window - enforce N per window
85+
--- pipelined_requests eval
86+
["GET /hello", "GET /hello", "GET /hello"]
87+
--- error_code eval
88+
[200, 200, 503]
89+
90+
91+
=== TEST 3: redis-cluster policy with sliding window - remaining header on reject
92+
--- config
93+
location /t {
94+
content_by_lua_block {
95+
local json = require "t.toolkit.json"
96+
local http = require "resty.http"
97+
local uri = "http://127.0.0.1:" .. ngx.var.server_port
98+
.. "/hello"
99+
local ress = {}
100+
101+
-- ensure previous windows are expired before starting this test
102+
ngx.sleep(2.2)
103+
104+
-- first request: allowed, remaining should be 1
105+
do
106+
local httpc = http.new()
107+
local res, err = httpc:request_uri(uri)
108+
if not res then
109+
ngx.say(err)
110+
return
111+
end
112+
table.insert(ress, {res.status, res.headers["X-RateLimit-Remaining"]})
113+
end
114+
115+
-- second request: allowed, remaining should be 0
116+
do
117+
local httpc = http.new()
118+
local res, err = httpc:request_uri(uri)
119+
if not res then
120+
ngx.say(err)
121+
return
122+
end
123+
table.insert(ress, {res.status, res.headers["X-RateLimit-Remaining"]})
124+
end
125+
126+
-- third request: rejected, remaining header should stay at 0
127+
do
128+
local httpc = http.new()
129+
local res, err = httpc:request_uri(uri)
130+
if not res then
131+
ngx.say(err)
132+
return
133+
end
134+
table.insert(ress, {res.status, res.headers["X-RateLimit-Remaining"]})
135+
end
136+
137+
ngx.say(json.encode(ress))
138+
}
139+
}
140+
--- response_body
141+
[[200,"1"],[200,"0"],[503,"0"]]
142+
143+
144+
=== TEST 4: redis-cluster policy with sliding window - allow after window passes
145+
--- config
146+
location /t {
147+
content_by_lua_block {
148+
local json = require "t.toolkit.json"
149+
local http = require "resty.http"
150+
local uri = "http://127.0.0.1:" .. ngx.var.server_port
151+
.. "/hello"
152+
local codes = {}
153+
154+
-- ensure previous windows are expired before starting this test
155+
ngx.sleep(2.2)
156+
157+
-- consume full quota
158+
for i = 1, 2 do
159+
local httpc = http.new()
160+
local res, err = httpc:request_uri(uri)
161+
if not res then
162+
ngx.say(err)
163+
return
164+
end
165+
table.insert(codes, res.status)
166+
end
167+
168+
-- wait longer than the sliding window (2s)
169+
ngx.sleep(2.2)
170+
171+
-- should be allowed again after window has passed
172+
do
173+
local httpc = http.new()
174+
local res, err = httpc:request_uri(uri)
175+
if not res then
176+
ngx.say(err)
177+
return
178+
end
179+
table.insert(codes, res.status)
180+
end
181+
182+
ngx.say(json.encode(codes))
183+
}
184+
}
185+
--- response_body
186+
[200,200,200]
187+
188+

0 commit comments

Comments
 (0)