Skip to content

Commit 1af04e0

Browse files
shreemaan-abhishekshreemaan-abhishek
committed
feat: add ttl to redis counters in redis policy rate limiters
Signed-off-by: Abhishek Choudhary <shreemaan.abhishek@gmail.com>
1 parent bcf0c04 commit 1af04e0

File tree

8 files changed

+753
-3
lines changed

8 files changed

+753
-3
lines changed

‎apisix/plugins/limit-conn.lua‎

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
local core = require("apisix.core")
1818
local limit_conn = require("apisix.plugins.limit-conn.init")
1919
local redis_schema = require("apisix.utils.redis-schema")
20-
local policy_to_additional_properties = redis_schema.schema
2120
local plugin_name = "limit-conn"
2221
local workflow = require("apisix.plugins.workflow")
2322

@@ -55,7 +54,7 @@ local schema = {
5554
},
5655
},
5756
},
58-
["then"] = policy_to_additional_properties.redis,
57+
["then"] = redis_schema.ttl_policy_schema("redis", 60),
5958
["else"] = {
6059
["if"] = {
6160
properties = {
@@ -64,7 +63,7 @@ local schema = {
6463
},
6564
},
6665
},
67-
["then"] = policy_to_additional_properties["redis-cluster"],
66+
["then"] = redis_schema.ttl_policy_schema("redis-cluster", 60),
6867
}
6968
}
7069

‎apisix/plugins/limit-conn/util.lua‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@ function _M.incoming(self, red, key, commit)
3333
return nil, err
3434
end
3535

36+
if self.conf.key_ttl then
37+
red:expire(key, self.conf.key_ttl)
38+
end
39+
3640
if conn > max + self.burst then
3741
conn, err = red:incrby(key, -1)
3842
if not conn then
@@ -69,6 +73,10 @@ function _M.leaving(self, red, key, req_latency)
6973
return nil, err
7074
end
7175

76+
if self.conf.key_ttl then
77+
red:expire(key, self.conf.key_ttl)
78+
end
79+
7280
if req_latency then
7381
local unit_delay = self.unit_delay
7482
self.unit_delay = (req_latency + unit_delay) / 2

‎apisix/plugins/limit-req/util.lua‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,10 @@ function _M.incoming(self, red, key, commit)
6868
if not ok then
6969
return nil, err
7070
end
71+
72+
local ttl = math.ceil(self.burst / self.rate) + 1
73+
red:expire(excess_key, ttl)
74+
red:expire(last_key, ttl)
7175
end
7276

7377
-- return the delay in seconds, as well as excess

‎apisix/utils/redis-schema.lua‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,4 +78,14 @@ local _M = {
7878
schema = policy_to_additional_properties
7979
}
8080

81+
82+
function _M.ttl_policy_schema(kind, ttl)
83+
local schema = policy_to_additional_properties[kind]
84+
schema.properties.key_ttl = {
85+
type = "integer", default = ttl,
86+
}
87+
return schema
88+
end
89+
90+
8191
return _M

‎t/plugin/limit-conn-redis-cluster-ttl.t‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,260 @@
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_root_location();
23+
no_shuffle();
24+
25+
add_block_preprocessor(sub {
26+
my ($block) = @_;
27+
my $config = $block->config // <<_EOC_;
28+
location /check {
29+
content_by_lua_block {
30+
local redis_cluster = require "resty.rediscluster"
31+
local config = {
32+
name = "test-cluster",
33+
serv_list = {
34+
{ ip = "127.0.0.1", port = 5000 },
35+
{ ip = "127.0.0.1", port = 5001 },
36+
{ ip = "127.0.0.1", port = 5002 },
37+
{ ip = "127.0.0.1", port = 5003 },
38+
{ ip = "127.0.0.1", port = 5004 },
39+
{ ip = "127.0.0.1", port = 5005 },
40+
},
41+
keepalive_timeout = 60000,
42+
keepalive_cons = 1000,
43+
connect_timeout = 1000,
44+
socket_timeout = 1000
45+
}
46+
local red = redis_cluster:new(config)
47+
48+
-- make a request to /access
49+
local httpc = require("resty.http").new()
50+
local uri = "http://127.0.0.1:" .. ngx.var.server_port .. "/access"
51+
local res, err = httpc:request_uri(uri, {
52+
method = "GET"
53+
})
54+
55+
if not res then
56+
ngx.say("failed to request: ", err)
57+
return
58+
end
59+
60+
local key = "limit_conn:127.0.0.1"
61+
local ttl, err = red:ttl(key)
62+
63+
if not ttl then
64+
ngx.say("failed to get ttl: ", err)
65+
return
66+
end
67+
68+
if ttl > 50 and ttl <= 60 then
69+
ngx.say("ttl is 60")
70+
else
71+
ngx.say("ttl is " .. tostring(ttl))
72+
end
73+
}
74+
}
75+
_EOC_
76+
$block->set_value("config", $config);
77+
});
78+
79+
run_tests;
80+
81+
__DATA__
82+
83+
=== TEST 1: check schema
84+
--- config
85+
location /t {
86+
content_by_lua_block {
87+
local plugin = require("apisix.plugins.limit-conn")
88+
local ok, err = plugin.check_schema({
89+
conn = 1,
90+
burst = 0,
91+
default_conn_delay = 0.1,
92+
rejected_code = 503,
93+
key = 'remote_addr',
94+
policy = "redis-cluster",
95+
redis_cluster_nodes = {
96+
"127.0.0.1:5000",
97+
"127.0.0.1:5001"
98+
},
99+
redis_cluster_name = "test",
100+
key_ttl = 10,
101+
})
102+
if not ok then
103+
ngx.say(err)
104+
end
105+
ngx.say("ok")
106+
}
107+
}
108+
--- request
109+
GET /t
110+
--- response_body
111+
ok
112+
113+
114+
115+
=== TEST 2: trigger limit-conn and check default TTL
116+
--- config
117+
location /test {
118+
content_by_lua_block {
119+
local t = require("lib.test_admin").test
120+
local code, body = t('/apisix/admin/routes/1',
121+
ngx.HTTP_PUT,
122+
[[{
123+
"plugins": {
124+
"limit-conn": {
125+
"conn": 100,
126+
"burst": 50,
127+
"default_conn_delay": 0.1,
128+
"key": "remote_addr",
129+
"policy": "redis-cluster",
130+
"redis_cluster_nodes": [
131+
"127.0.0.1:5000",
132+
"127.0.0.1:5001",
133+
"127.0.0.1:5002",
134+
"127.0.0.1:5003",
135+
"127.0.0.1:5004",
136+
"127.0.0.1:5005"
137+
],
138+
"redis_cluster_name": "test-cluster"
139+
}
140+
},
141+
"upstream": {
142+
"nodes": {
143+
"127.0.0.1:1980": 1
144+
},
145+
"type": "roundrobin"
146+
},
147+
"uri": "/access"
148+
}]]
149+
)
150+
ngx.say(body)
151+
}
152+
}
153+
--- request
154+
GET /test
155+
--- response_body
156+
passed
157+
158+
159+
160+
=== TEST 3: access and check default ttl (should be 60)
161+
--- request
162+
GET /check
163+
--- response_body
164+
ttl is 60
165+
166+
167+
168+
=== TEST 4: configure custom TTL
169+
--- config
170+
location /test_update {
171+
content_by_lua_block {
172+
local t = require("lib.test_admin").test
173+
local code, body = t('/apisix/admin/routes/1',
174+
ngx.HTTP_PUT,
175+
[[{
176+
"plugins": {
177+
"limit-conn": {
178+
"conn": 100,
179+
"burst": 50,
180+
"default_conn_delay": 0.1,
181+
"key": "remote_addr",
182+
"policy": "redis-cluster",
183+
"redis_cluster_nodes": [
184+
"127.0.0.1:5000",
185+
"127.0.0.1:5001",
186+
"127.0.0.1:5002",
187+
"127.0.0.1:5003",
188+
"127.0.0.1:5004",
189+
"127.0.0.1:5005"
190+
],
191+
"redis_cluster_name": "test-cluster",
192+
"key_ttl": 10
193+
}
194+
},
195+
"upstream": {
196+
"nodes": {
197+
"127.0.0.1:1980": 1
198+
},
199+
"type": "roundrobin"
200+
},
201+
"uri": "/access"
202+
}]]
203+
)
204+
ngx.say(body)
205+
}
206+
}
207+
--- request
208+
GET /test_update
209+
--- response_body
210+
passed
211+
212+
213+
214+
=== TEST 5: access and check custom ttl (should be 10)
215+
--- config
216+
location /check_custom {
217+
content_by_lua_block {
218+
local redis_cluster = require "resty.rediscluster"
219+
local config = {
220+
name = "test-cluster",
221+
serv_list = {
222+
{ ip = "127.0.0.1", port = 5000 },
223+
{ ip = "127.0.0.1", port = 5001 },
224+
{ ip = "127.0.0.1", port = 5002 },
225+
{ ip = "127.0.0.1", port = 5003 },
226+
{ ip = "127.0.0.1", port = 5004 },
227+
{ ip = "127.0.0.1", port = 5005 },
228+
},
229+
keepalive_timeout = 60000,
230+
keepalive_cons = 1000,
231+
connect_timeout = 1000,
232+
socket_timeout = 1000
233+
}
234+
local red = redis_cluster:new(config)
235+
236+
local httpc = require("resty.http").new()
237+
local uri = "http://127.0.0.1:" .. ngx.var.server_port .. "/access"
238+
local res, err = httpc:request_uri(uri, {
239+
method = "GET"
240+
})
241+
242+
local key = "limit_conn:127.0.0.1"
243+
local ttl, err = red:ttl(key)
244+
245+
if not ttl then
246+
ngx.say("failed to get ttl: ", err)
247+
return
248+
end
249+
250+
if ttl > 5 and ttl <= 10 then
251+
ngx.say("ttl is 10")
252+
else
253+
ngx.say("ttl is " .. tostring(ttl))
254+
end
255+
}
256+
}
257+
--- request
258+
GET /check_custom
259+
--- response_body
260+
ttl is 10

0 commit comments

Comments
 (0)
⚡