feat: support configuring variables in limit-conn, limit-count and ai-rate-limiting (#12967)

Author: shreemaan-abhishek

apisix/plugins/ai-rate-limiting.lua

@@ -27,17 +27,37 @@ local instance_limit_schema = {
     type = "object",
     properties = {
         name = {type = "string"},
-        limit = {type = "integer", minimum = 1},
-        time_window = {type = "integer", minimum = 1}
+        limit = {
+            oneOf = {
+                {type = "integer", minimum = 1},
+                {type = "string"},
+            },
+        },
+        time_window = {
+            oneOf = {
+                {type = "integer", minimum = 1},
+                {type = "string"},
+            },
+        }
     },
     required = {"name", "limit", "time_window"}
 }
 
 local schema = {
     type = "object",
     properties = {
-        limit = {type = "integer", exclusiveMinimum = 0},
-        time_window = {type = "integer",  exclusiveMinimum = 0},
+        limit = {
+            oneOf = {
+                {type = "integer", exclusiveMinimum = 0},
+                {type = "string"},
+            },
+        },
+        time_window = {
+            oneOf = {
+                {type = "integer", exclusiveMinimum = 0},
+                {type = "string"},
+            },
+        },
         show_limit_quota_header = {type = "boolean", default = true},
         limit_strategy = {
             type = "string",

apisix/plugins/limit-conn.lua

@@ -24,8 +24,18 @@ local workflow                           = require("apisix.plugins.workflow")
 local schema = {
     type = "object",
     properties = {
-        conn = {type = "integer", exclusiveMinimum = 0},               -- limit.conn max
-        burst = {type = "integer",  minimum = 0},
+        conn = {
+            oneOf = {
+                {type = "integer", exclusiveMinimum = 0},
+                {type = "string"},
+            },
+        },
+        burst = {
+            oneOf = {
+                {type = "integer", minimum = 0},
+                {type = "string"},
+            },
+        },
         default_conn_delay = {type = "number", exclusiveMinimum = 0},
         only_use_default_delay = {type = "boolean", default = false},
         key = {type = "string"},

apisix/plugins/limit-conn/init.lua

@@ -18,6 +18,9 @@ local limit_conn_new = require("resty.limit.conn").new
 local core = require("apisix.core")
 local is_http = ngx.config.subsystem == "http"
 local sleep = core.sleep
+local tonumber = tonumber
+local type = type
+local tostring = tostring
 local shdict_name = "plugin-limit-conn"
 if ngx.config.subsystem == "stream" then
     shdict_name = shdict_name .. "-stream"
@@ -34,38 +37,63 @@ do
 end
 
 
-local lrucache = core.lrucache.new({
-    type = "plugin",
-})
 local _M = {}
 
 
-local function create_limit_obj(conf)
+local function create_limit_obj(ctx, conf)
     core.log.info("create new limit-conn plugin instance")
 
+    local conn = conf.conn
+    if type(conn) == "string" then
+        local err, _
+        conn, err, _ = core.utils.resolve_var(conn, ctx.var)
+        if err then
+            return nil, "could not resolve vars in conn: " .. err
+        end
+        conn = tonumber(conn)
+        if not conn then
+            return nil, "resolved conn is not a number: " .. tostring(conn)
+        end
+    end
+
+    local burst = conf.burst
+    if type(burst) == "string" then
+        local err, _
+        burst, err, _ = core.utils.resolve_var(burst, ctx.var)
+        if err then
+            return nil, "could not resolve vars in burst: " .. err
+        end
+        burst = tonumber(burst)
+        if not burst then
+            return nil, "resolved burst is not a number: " .. tostring(burst)
+        end
+    end
+
+    core.log.info("limit conn: ", conn, ", burst: ", burst)
+
     if conf.policy == "redis" then
         core.log.info("create new limit-conn redis plugin instance")
 
-        return redis_single_new("plugin-limit-conn", conf, conf.conn, conf.burst,
+        return redis_single_new("plugin-limit-conn", conf, conn, burst,
                                 conf.default_conn_delay)
 
     elseif conf.policy == "redis-cluster" then
 
         core.log.info("create new limit-conn redis-cluster plugin instance")
 
-        return redis_cluster_new("plugin-limit-conn", conf, conf.conn, conf.burst,
+        return redis_cluster_new("plugin-limit-conn", conf, conn, burst,
                                  conf.default_conn_delay)
     else
         core.log.info("create new limit-conn plugin instance")
-        return limit_conn_new(shdict_name, conf.conn, conf.burst,
+        return limit_conn_new(shdict_name, conn, burst,
                               conf.default_conn_delay)
     end
 end
 
 
 function _M.increase(conf, ctx)
     core.log.info("ver: ", ctx.conf_version)
-    local lim, err = lrucache(conf, nil, create_limit_obj, conf)
+    local lim, err = create_limit_obj(ctx, conf)
     if not lim then
         core.log.error("failed to instantiate a resty.limit.conn object: ", err)
         if conf.allow_degradation then

apisix/plugins/limit-count/init.lua

@@ -22,6 +22,9 @@ local pairs = pairs
 local redis_schema = require("apisix.utils.redis-schema")
 local policy_to_additional_properties = redis_schema.schema
 local get_phase = ngx.get_phase
+local tonumber = tonumber
+local type = type
+local tostring = tostring
 
 local limit_redis_cluster_new
 local limit_redis_new
@@ -36,9 +39,6 @@ do
     local cluster_src = "apisix.plugins.limit-count.limit-count-redis-cluster"
     limit_redis_cluster_new = require(cluster_src).new
 end
-local lrucache = core.lrucache.new({
-    type = 'plugin', serial_creating = true,
-})
 local group_conf_lru = core.lrucache.new({
     type = 'plugin',
 })
@@ -70,8 +70,18 @@ local metadata_schema = {
 local schema = {
     type = "object",
     properties = {
-        count = {type = "integer", exclusiveMinimum = 0},
-        time_window = {type = "integer",  exclusiveMinimum = 0},
+        count = {
+            oneOf = {
+                {type = "integer", exclusiveMinimum = 0},
+                {type = "string"},
+            },
+        },
+        time_window = {
+            oneOf = {
+                {type = "integer", exclusiveMinimum = 0},
+                {type = "string"},
+            },
+        },
         group = {type = "string"},
         key = {type = "string", default = "remote_addr"},
         key_type = {type = "string",
@@ -174,22 +184,47 @@ function _M.check_schema(conf, schema_type)
 end
 
 
-local function create_limit_obj(conf, plugin_name)
+local function create_limit_obj(conf, ctx, plugin_name)
     core.log.info("create new " .. plugin_name .. " plugin instance")
 
+    local count = conf.count
+    if type(count) == "string" then
+        local err, _
+        count, err, _ = core.utils.resolve_var(count, ctx.var)
+        if err then
+            return nil, "could not resolve vars in count: " .. err
+        end
+        count = tonumber(count)
+        if not count then
+            return nil, "resolved count is not a number: " .. tostring(count)
+        end
+    end
+
+    local time_window = conf.time_window
+    if type(time_window) == "string" then
+        local err, _
+        time_window, err, _ = core.utils.resolve_var(time_window, ctx.var)
+        if err then
+            return nil, "could not resolve vars in time_window: " .. err
+        end
+        time_window = tonumber(time_window)
+        if not time_window then
+            return nil, "resolved time_window is not a number: " .. tostring(time_window)
+        end
+    end
+
+    core.log.info("limit count: ", count, ", time_window: ", time_window)
+
     if not conf.policy or conf.policy == "local" then
-        return limit_local_new("plugin-" .. plugin_name, conf.count,
-                               conf.time_window)
+        return limit_local_new("plugin-" .. plugin_name, count, time_window)
     end
 
     if conf.policy == "redis" then
-        return limit_redis_new("plugin-" .. plugin_name,
-                               conf.count, conf.time_window, conf)
+        return limit_redis_new("plugin-" .. plugin_name, count, time_window, conf)
     end
 
     if conf.policy == "redis-cluster" then
-        return limit_redis_cluster_new("plugin-" .. plugin_name, conf.count,
-                                       conf.time_window, conf)
+        return limit_redis_cluster_new("plugin-" .. plugin_name, count, time_window, conf)
     end
 
     return nil
@@ -223,26 +258,11 @@ local function gen_limit_key(conf, ctx, key)
 end
 
 
-local function gen_limit_obj(conf, ctx, plugin_name)
-    if conf.group then
-        return lrucache(conf.group, "", create_limit_obj, conf, plugin_name)
-    end
-
-    local extra_key
-    if conf._vid then
-        extra_key = conf.policy .. '#' .. conf._vid
-    else
-        extra_key = conf.policy
-    end
-
-    return core.lrucache.plugin_ctx(lrucache, ctx, extra_key, create_limit_obj, conf, plugin_name)
-end
-
 function _M.rate_limit(conf, ctx, name, cost, dry_run)
     core.log.info("ver: ", ctx.conf_version)
     core.log.info("conf: ", core.json.delay_encode(conf, true))
 
-    local lim, err = gen_limit_obj(conf, ctx, name)
+    local lim, err = create_limit_obj(conf, ctx, name)
 
     if not lim then
         core.log.error("failed to fetch limit.count object: ", err)
@@ -307,7 +327,7 @@ function _M.rate_limit(conf, ctx, name, cost, dry_run)
         if err == "rejected" then
             -- show count limit header when rejected
             if conf.show_limit_quota_header and set_header then
-                core.response.set_header(set_limit_headers.limit_header, conf.count,
+                core.response.set_header(set_limit_headers.limit_header, lim.limit,
                 set_limit_headers.remaining_header, 0,
                 set_limit_headers.reset_header, reset)
             end
@@ -326,7 +346,7 @@ function _M.rate_limit(conf, ctx, name, cost, dry_run)
     end
 
     if conf.show_limit_quota_header and set_header then
-        core.response.set_header(set_limit_headers.limit_header, conf.count,
+        core.response.set_header(set_limit_headers.limit_header, lim.limit,
             set_limit_headers.remaining_header, remaining,
             set_limit_headers.reset_header, reset)
     end

apisix/plugins/limit-count/limit-count-local.lua

@@ -57,7 +57,9 @@ function _M.new(plugin_name, limit, window)
 
     local self = {
         limit_count = limit_count.new(plugin_name, limit, window),
-        dict = ngx.shared[plugin_name .. "-reset-header"]
+        dict = ngx.shared[plugin_name .. "-reset-header"],
+        limit = limit,
+        window = window,
     }
 
     return setmetatable(self, mt)
@@ -67,8 +69,8 @@ function _M.incoming(self, key, commit, conf, cost)
     local delay, remaining = self.limit_count:incoming(key, commit, cost)
     local reset
 
-    if remaining == conf.count - cost then
-        reset = set_endtime(self, key, conf.time_window)
+    if remaining == self.limit - cost then
+        reset = set_endtime(self, key, self.window)
     else
         reset = read_reset(self, key)
     end