You are viewing a plain text version of this content. The canonical link for it is here.
Posted to notifications@apisix.apache.org by le...@apache.org on 2023/04/21 03:00:06 UTC
[apisix] branch master updated: feat(limit-count): support consuming specified cost in the limit count function (#9291)
This is an automated email from the ASF dual-hosted git repository.
leslie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/apisix.git
The following commit(s) were added to refs/heads/master by this push:
new fc6cb9818 feat(limit-count): support consuming specified cost in the limit count function (#9291)
fc6cb9818 is described below
commit fc6cb98187ee96bf9539ef1d77ea5287a89c7e7a
Author: Abhishek Choudhary <sh...@gmail.com>
AuthorDate: Fri Apr 21 08:29:57 2023 +0530
feat(limit-count): support consuming specified cost in the limit count function (#9291)
---
apisix/plugins/limit-count.lua | 2 +-
apisix/plugins/limit-count/init.lua | 23 +++---
apisix/plugins/limit-count/limit-count-local.lua | 13 ++--
.../limit-count/limit-count-redis-cluster.lua | 11 +--
apisix/plugins/limit-count/limit-count-redis.lua | 11 +--
apisix/plugins/workflow.lua | 2 +-
t/plugin/limit-count-redis-cluster3.t | 79 +++++++++++++++++++
t/plugin/limit-count-redis4.t | 87 +++++++++++++++++++++
t/plugin/limit-count5.t | 91 ++++++++++++++++++++++
9 files changed, 292 insertions(+), 27 deletions(-)
diff --git a/apisix/plugins/limit-count.lua b/apisix/plugins/limit-count.lua
index 0eafd6423..72ab574df 100644
--- a/apisix/plugins/limit-count.lua
+++ b/apisix/plugins/limit-count.lua
@@ -31,7 +31,7 @@ end
function _M.access(conf, ctx)
- return limit_count.rate_limit(conf, ctx)
+ return limit_count.rate_limit(conf, ctx, plugin_name, 1)
end
diff --git a/apisix/plugins/limit-count/init.lua b/apisix/plugins/limit-count/init.lua
index 40b18c17a..069ab6739 100644
--- a/apisix/plugins/limit-count/init.lua
+++ b/apisix/plugins/limit-count/init.lua
@@ -20,7 +20,6 @@ local tab_insert = table.insert
local ipairs = ipairs
local pairs = pairs
-local plugin_name = "limit-count"
local limit_redis_cluster_new
local limit_redis_new
local limit_local_new
@@ -200,8 +199,8 @@ function _M.check_schema(conf)
end
-local function create_limit_obj(conf)
- core.log.info("create new limit-count plugin instance")
+local function create_limit_obj(conf, plugin_name)
+ core.log.info("create new " .. plugin_name .. " plugin instance")
if not conf.policy or conf.policy == "local" then
return limit_local_new("plugin-" .. plugin_name, conf.count,
@@ -243,9 +242,9 @@ local function gen_limit_key(conf, ctx, key)
end
-local function gen_limit_obj(conf, ctx)
+local function gen_limit_obj(conf, ctx, plugin_name)
if conf.group then
- return lrucache(conf.group, "", create_limit_obj, conf)
+ return lrucache(conf.group, "", create_limit_obj, conf, plugin_name)
end
local extra_key
@@ -255,13 +254,13 @@ local function gen_limit_obj(conf, ctx)
extra_key = conf.policy
end
- return core.lrucache.plugin_ctx(lrucache, ctx, extra_key, create_limit_obj, conf)
+ return core.lrucache.plugin_ctx(lrucache, ctx, extra_key, create_limit_obj, conf, plugin_name)
end
-function _M.rate_limit(conf, ctx)
+function _M.rate_limit(conf, ctx, name, cost)
core.log.info("ver: ", ctx.conf_version)
- local lim, err = gen_limit_obj(conf, ctx)
+ local lim, err = gen_limit_obj(conf, ctx, name)
if not lim then
core.log.error("failed to fetch limit.count object: ", err)
@@ -298,7 +297,13 @@ function _M.rate_limit(conf, ctx)
key = gen_limit_key(conf, ctx, key)
core.log.info("limit key: ", key)
- local delay, remaining, reset = lim:incoming(key, true, conf)
+ local delay, remaining, reset
+ if not conf.policy or conf.policy == "local" then
+ delay, remaining, reset = lim:incoming(key, true, conf, cost)
+ else
+ delay, remaining, reset = lim:incoming(key, cost)
+ end
+
if not delay then
local err = remaining
if err == "rejected" then
diff --git a/apisix/plugins/limit-count/limit-count-local.lua b/apisix/plugins/limit-count/limit-count-local.lua
index f69094031..27c1fc454 100644
--- a/apisix/plugins/limit-count/limit-count-local.lua
+++ b/apisix/plugins/limit-count/limit-count-local.lua
@@ -14,7 +14,8 @@
-- See the License for the specific language governing permissions and
-- limitations under the License.
--
-local limit_local_new = require("resty.limit.count").new
+local limit_count = require("resty.limit.count")
+
local ngx = ngx
local ngx_time = ngx.time
local assert = assert
@@ -55,21 +56,21 @@ function _M.new(plugin_name, limit, window)
assert(limit > 0 and window > 0)
local self = {
- limit_count = limit_local_new(plugin_name, limit, window),
- dict = ngx.shared["plugin-limit-count-reset-header"]
+ limit_count = limit_count.new(plugin_name, limit, window),
+ dict = ngx.shared[plugin_name .. "-reset-header"]
}
return setmetatable(self, mt)
end
-function _M.incoming(self, key, commit, conf)
- local delay, remaining = self.limit_count:incoming(key, commit)
+function _M.incoming(self, key, commit, conf, cost)
+ local delay, remaining = self.limit_count:incoming(key, commit, cost)
local reset = 0
if not delay then
return delay, remaining, reset
end
- if remaining == conf.count - 1 then
+ if remaining == conf.count - cost then
reset = set_endtime(self, key, conf.time_window)
else
reset = read_reset(self, key)
diff --git a/apisix/plugins/limit-count/limit-count-redis-cluster.lua b/apisix/plugins/limit-count/limit-count-redis-cluster.lua
index 25a3c1c57..7800f1c88 100644
--- a/apisix/plugins/limit-count/limit-count-redis-cluster.lua
+++ b/apisix/plugins/limit-count/limit-count-redis-cluster.lua
@@ -30,12 +30,13 @@ local mt = {
local script = core.string.compress_script([=[
+ assert(tonumber(ARGV[3]) >= 1, "cost must be at least 1")
local ttl = redis.call('ttl', KEYS[1])
if ttl < 0 then
- redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
- return {ARGV[1] - 1, ARGV[2]}
+ redis.call('set', KEYS[1], ARGV[1] - ARGV[3], 'EX', ARGV[2])
+ return {ARGV[1] - ARGV[3], ARGV[2]}
end
- return {redis.call('incrby', KEYS[1], -1), ttl}
+ return {redis.call('incrby', KEYS[1], 0 - ARGV[3]), ttl}
]=])
@@ -90,14 +91,14 @@ function _M.new(plugin_name, limit, window, conf)
end
-function _M.incoming(self, key)
+function _M.incoming(self, key, cost)
local red = self.red_cli
local limit = self.limit
local window = self.window
key = self.plugin_name .. tostring(key)
local ttl = 0
- local res, err = red:eval(script, 1, key, limit, window)
+ local res, err = red:eval(script, 1, key, limit, window, cost or 1)
if err then
return nil, err, ttl
diff --git a/apisix/plugins/limit-count/limit-count-redis.lua b/apisix/plugins/limit-count/limit-count-redis.lua
index f2e441d61..85f769cbd 100644
--- a/apisix/plugins/limit-count/limit-count-redis.lua
+++ b/apisix/plugins/limit-count/limit-count-redis.lua
@@ -30,12 +30,13 @@ local mt = {
local script = core.string.compress_script([=[
+ assert(tonumber(ARGV[3]) >= 1, "cost must be at least 1")
local ttl = redis.call('ttl', KEYS[1])
if ttl < 0 then
- redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
- return {ARGV[1] - 1, ARGV[2]}
+ redis.call('set', KEYS[1], ARGV[1] - ARGV[3], 'EX', ARGV[2])
+ return {ARGV[1] - ARGV[3], ARGV[2]}
end
- return {redis.call('incrby', KEYS[1], -1), ttl}
+ return {redis.call('incrby', KEYS[1], 0 - ARGV[3]), ttl}
]=])
local function redis_cli(conf)
@@ -95,7 +96,7 @@ function _M.new(plugin_name, limit, window, conf)
return setmetatable(self, mt)
end
-function _M.incoming(self, key)
+function _M.incoming(self, key, cost)
local conf = self.conf
local red, err = redis_cli(conf)
if not red then
@@ -108,7 +109,7 @@ function _M.incoming(self, key)
key = self.plugin_name .. tostring(key)
local ttl = 0
- res, err = red:eval(script, 1, key, limit, window)
+ res, err = red:eval(script, 1, key, limit, window, cost or 1)
if err then
return nil, err, ttl
diff --git a/apisix/plugins/workflow.lua b/apisix/plugins/workflow.lua
index 5adb5b455..4a8153f43 100644
--- a/apisix/plugins/workflow.lua
+++ b/apisix/plugins/workflow.lua
@@ -93,7 +93,7 @@ end
local function rate_limit(conf, ctx)
- return limit_count.rate_limit(conf, ctx)
+ return limit_count.rate_limit(conf, ctx, "limit-count", 1)
end
diff --git a/t/plugin/limit-count-redis-cluster3.t b/t/plugin/limit-count-redis-cluster3.t
new file mode 100644
index 000000000..b534c3a2c
--- /dev/null
+++ b/t/plugin/limit-count-redis-cluster3.t
@@ -0,0 +1,79 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+use t::APISIX;
+
+my $nginx_binary = $ENV{'TEST_NGINX_BINARY'} || 'nginx';
+my $version = eval { `$nginx_binary -V 2>&1` };
+
+if ($version !~ m/\/apisix-nginx-module/) {
+ plan(skip_all => "apisix-nginx-module not installed");
+} else {
+ plan('no_plan');
+}
+
+repeat_each(1);
+no_long_string();
+no_shuffle();
+no_root_location();
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!$block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ if (!$block->error_log && !$block->no_error_log) {
+ $block->set_value("no_error_log", "[error]\n[alert]");
+ }
+});
+
+run_tests;
+
+__DATA__
+
+=== TEST 1: modified redis script, cost == 2
+--- config
+ location /t {
+ content_by_lua_block {
+ local conf = {
+ redis_cluster_nodes = {"127.0.0.1:5000", "127.0.0.1:5001"},
+ redis_cluster_name = "redis-cluster-1",
+ redis_cluster_ssl = false,
+ redis_timeout = 1000,
+ key_type = "var",
+ time_window = 60,
+ show_limit_quota_header = true,
+ allow_degradation = false,
+ key = "remote_addr",
+ rejected_code = 503,
+ count = 3,
+ policy = "redis-cluster",
+ redis_cluster_ssl_verify = false
+ }
+
+ local lim_count_redis_cluster = require("apisix.plugins.limit-count.limit-count-redis-cluster")
+ local lim = lim_count_redis_cluster.new("limit-count", 3, 60, conf)
+ local uri = ngx.var.uri
+ local _, remaining, _ = lim:incoming(uri, 2)
+
+ ngx.say("remaining: ", remaining)
+ }
+ }
+--- response_body
+remaining: 1
diff --git a/t/plugin/limit-count-redis4.t b/t/plugin/limit-count-redis4.t
new file mode 100644
index 000000000..47deedf06
--- /dev/null
+++ b/t/plugin/limit-count-redis4.t
@@ -0,0 +1,87 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+BEGIN {
+ if ($ENV{TEST_NGINX_CHECK_LEAK}) {
+ $SkipReason = "unavailable for the hup tests";
+
+ } else {
+ $ENV{TEST_NGINX_USE_HUP} = 1;
+ undef $ENV{TEST_NGINX_USE_STAP};
+ }
+}
+
+use t::APISIX;
+
+my $nginx_binary = $ENV{'TEST_NGINX_BINARY'} || 'nginx';
+my $version = eval { `$nginx_binary -V 2>&1` };
+
+if ($version !~ m/\/apisix-nginx-module/) {
+ plan(skip_all => "apisix-nginx-module not installed");
+} else {
+ plan('no_plan');
+}
+
+repeat_each(1);
+no_long_string();
+no_shuffle();
+no_root_location();
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!$block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ if (!$block->error_log && !$block->no_error_log) {
+ $block->set_value("no_error_log", "[error]\n[alert]");
+ }
+});
+
+run_tests;
+
+__DATA__
+
+=== TEST 1: modified redis script, cost == 2
+--- config
+ location /t {
+ content_by_lua_block {
+ local conf = {
+ allow_degradation = false,
+ rejected_code = 503,
+ redis_timeout = 1000,
+ key_type = "var",
+ time_window = 60,
+ show_limit_quota_header = true,
+ count = 3,
+ redis_host = "127.0.0.1",
+ redis_port = 6379,
+ redis_database = 0,
+ policy = "redis",
+ key = "remote_addr"
+ }
+
+ local lim_count_redis = require("apisix.plugins.limit-count.limit-count-redis")
+ local lim = lim_count_redis.new("limit-count", 3, 60, conf)
+ local uri = ngx.var.uri
+ local _, remaining, _ = lim:incoming(uri, 2)
+
+ ngx.say("remaining: ", remaining)
+ }
+ }
+--- response_body
+remaining: 1
diff --git a/t/plugin/limit-count5.t b/t/plugin/limit-count5.t
new file mode 100644
index 000000000..47daaa71c
--- /dev/null
+++ b/t/plugin/limit-count5.t
@@ -0,0 +1,91 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+BEGIN {
+ if ($ENV{TEST_NGINX_CHECK_LEAK}) {
+ $SkipReason = "unavailable for the hup tests";
+
+ } else {
+ $ENV{TEST_NGINX_USE_HUP} = 1;
+ undef $ENV{TEST_NGINX_USE_STAP};
+ }
+}
+
+use t::APISIX;
+
+my $nginx_binary = $ENV{'TEST_NGINX_BINARY'} || 'nginx';
+my $version = eval { `$nginx_binary -V 2>&1` };
+
+if ($version !~ m/\/apisix-nginx-module/) {
+ plan(skip_all => "apisix-nginx-module not installed");
+} else {
+ plan('no_plan');
+}
+
+add_block_preprocessor(sub {
+ my ($block) = @_;
+
+ if (!$block->request) {
+ $block->set_value("request", "GET /t");
+ }
+
+ if (!$block->error_log && !$block->no_error_log) {
+ $block->set_value("no_error_log", "[error]\n[alert]");
+ }
+});
+
+run_tests;
+
+__DATA__
+
+=== TEST 1: modified limit-count.incoming, cost == 2
+--- config
+ location = /t {
+ content_by_lua_block {
+ local conf = {
+ time_window = 60,
+ count = 10,
+ allow_degradation = false,
+ key_type = "var",
+ policy = "local",
+ rejected_code = 503,
+ show_limit_quota_header = true,
+ key = "remote_addr"
+ }
+ local limit_count_local = require "apisix.plugins.limit-count.limit-count-local"
+ local lim = limit_count_local.new("plugin-limit-count", 10, 60)
+ local uri = ngx.var.uri
+ for i = 1, 7 do
+ local delay, err = lim:incoming(uri, true, conf, 2)
+ if not delay then
+ ngx.say(err)
+ else
+ local remaining = err
+ ngx.say("remaining: ", remaining)
+ end
+ end
+ }
+ }
+--- request
+ GET /t
+--- response_body
+remaining: 8
+remaining: 6
+remaining: 4
+remaining: 2
+remaining: 0
+rejected
+rejected