fix(limit-conn): implement configurable redis key expiry (#12872)

Author: shreemaan-abhishek

apisix/plugins/limit-conn.lua

@@ -17,7 +17,6 @@
 local core                              = require("apisix.core")
 local limit_conn                        = require("apisix.plugins.limit-conn.init")
 local redis_schema                      = require("apisix.utils.redis-schema")
-local policy_to_additional_properties   = redis_schema.schema
 local plugin_name                       = "limit-conn"
 local workflow                           = require("apisix.plugins.workflow")
 
@@ -55,7 +54,7 @@ local schema = {
             },
         },
     },
-    ["then"] = policy_to_additional_properties.redis,
+    ["then"] = redis_schema.limit_conn_redis_schema,
     ["else"] = {
         ["if"] = {
             properties = {
@@ -64,7 +63,7 @@ local schema = {
                 },
             },
         },
-        ["then"] = policy_to_additional_properties["redis-cluster"],
+        ["then"] = redis_schema.limit_conn_redis_cluster_schema,
     }
 }
 

apisix/plugins/limit-conn/limit-conn-redis-cluster.lua

@@ -19,6 +19,7 @@ local core              = require("apisix.core")
 local util              = require("apisix.plugins.limit-conn.util")
 local setmetatable      = setmetatable
 local ngx_timer_at      = ngx.timer.at
+local ngx               = ngx
 
 local _M = {version = 0.1}
 
@@ -41,6 +42,7 @@ function _M.new(plugin_name, conf, max, burst, default_conn_delay)
         max = max + 0,    -- just to ensure the param is good
         unit_delay = default_conn_delay,
         red_cli = red_cli,
+        use_evalsha = false,
     }
     return setmetatable(self, mt)
 end
@@ -56,14 +58,19 @@ function _M.is_committed(self)
 end
 
 
-local function leaving_thread(premature, self, key, req_latency)
-    return util.leaving(self, self.red_cli, key, req_latency)
+local function leaving_thread(premature, self, key, req_latency, req_id)
+    return util.leaving(self, self.red_cli, key, req_latency, req_id)
 end
 
 
 function _M.leaving(self, key, req_latency)
+    local req_id
+    if ngx.ctx.limit_conn_req_ids then
+        req_id = ngx.ctx.limit_conn_req_ids[key]
+    end
+
     -- log_by_lua can't use cosocket
-    local ok, err = ngx_timer_at(0, leaving_thread, self, key, req_latency)
+    local ok, err = ngx_timer_at(0, leaving_thread, self, key, req_latency, req_id)
     if not ok then
         core.log.error("failed to create timer: ", err)
         return nil, err

apisix/plugins/limit-conn/limit-conn-redis.lua

@@ -18,6 +18,7 @@ local redis         = require("apisix.utils.redis")
 local core          = require("apisix.core")
 local util          = require("apisix.plugins.limit-conn.util")
 local ngx_timer_at  = ngx.timer.at
+local ngx           = ngx
 
 local setmetatable  = setmetatable
 
@@ -37,6 +38,7 @@ function _M.new(plugin_name, conf, max, burst, default_conn_delay)
         burst = burst,
         max = max + 0,    -- just to ensure the param is good
         unit_delay = default_conn_delay,
+        use_evalsha = true,
     }
     return setmetatable(self, mt)
 end
@@ -57,20 +59,25 @@ function _M.is_committed(self)
 end
 
 
-local function leaving_thread(premature, self, key, req_latency)
+local function leaving_thread(premature, self, key, req_latency, req_id)
 
     local conf = self.conf
     local red, err = redis.new(conf)
     if not red then
         return red, err
     end
-    return util.leaving(self, red, key, req_latency)
+    return util.leaving(self, red, key, req_latency, req_id)
 end
 
 
 function _M.leaving(self, key, req_latency)
+    local req_id
+    if ngx.ctx.limit_conn_req_ids then
+        req_id = ngx.ctx.limit_conn_req_ids[key]
+    end
+
     -- log_by_lua can't use cosocket
-    local ok, err = ngx_timer_at(0, leaving_thread, self, key, req_latency)
+    local ok, err = ngx_timer_at(0, leaving_thread, self, key, req_latency, req_id)
     if not ok then
         core.log.error("failed to create timer: ", err)
         return nil, err

apisix/plugins/limit-conn/util.lua

@@ -18,36 +18,100 @@
 local assert            = assert
 local math              = require "math"
 local floor             = math.floor
+local ngx               = ngx
+local ngx_time          = ngx.time
+local uuid              = require("resty.jit-uuid")
+local core              = require("apisix.core")
+
 local _M = {version = 0.3}
+local redis_incoming_script = core.string.compress_script([=[
+    local key = KEYS[1]
+    local limit = tonumber(ARGV[1])
+    local ttl = tonumber(ARGV[2])
+    local now = tonumber(ARGV[3])
+    local req_id = ARGV[4]
+
+    redis.call('ZREMRANGEBYSCORE', key, 0, now)
+
+    local count = redis.call('ZCARD', key)
+    if count >= limit then
+        return {0, count}
+    end
+
+    redis.call('ZADD', key, now + ttl, req_id)
+    redis.call('EXPIRE', key, ttl)
+    return {1, count + 1}
+]=])
+local redis_incoming_script_sha
+
+
+local function generate_redis_sha1(red)
+    local sha1, err = red:script("LOAD", redis_incoming_script)
+    if not sha1 then
+        return nil, err
+    end
+    return sha1
+end
 
 
 function _M.incoming(self, red, key, commit)
     local max = self.max
     self.committed = false
+    local raw_key = key
     key = "limit_conn" .. ":" .. key
 
-    local conn, err
+    local conn
     if commit then
-        conn, err = red:incrby(key, 1)
-        if not conn then
-            return nil, err
+        local req_id = ngx.ctx.request_id or uuid.generate_v4()
+        if not ngx.ctx.limit_conn_req_ids then
+            ngx.ctx.limit_conn_req_ids = {}
         end
+        ngx.ctx.limit_conn_req_ids[raw_key] = req_id
+
+        local now = ngx_time()
+        local res, err
+
+        if self.use_evalsha then
+            if not redis_incoming_script_sha then
+                redis_incoming_script_sha, err = generate_redis_sha1(red)
+                if not redis_incoming_script_sha then
+                    core.log.error("failed to generate redis sha1: ", err)
+                    return nil, err
+                end
+            end
 
-        if conn > max + self.burst then
-            conn, err = red:incrby(key, -1)
-            if not conn then
-                return nil, err
+            res, err = red:evalsha(redis_incoming_script_sha, 1, key,
+                                   max + self.burst, self.conf.key_ttl, now, req_id)
+
+            if err and core.string.has_prefix(err, "NOSCRIPT") then
+                core.log.warn("redis evalsha failed: ", err, ". Falling back to eval...")
+                redis_incoming_script_sha = nil
+                res, err = red:eval(redis_incoming_script, 1, key,
+                                    max + self.burst, self.conf.key_ttl, now, req_id)
             end
+        else
+            res, err = red:eval(redis_incoming_script, 1, key,
+                                max + self.burst, self.conf.key_ttl, now, req_id)
+        end
+
+        if not res then
+            return nil, err
+        end
+
+        local allowed = res[1]
+        conn = res[2]
+
+        if allowed == 0 then
             return nil, "rejected"
         end
+
         self.committed = true
 
     else
-        local conn_from_red, err = red:get(key)
-        if err then
-            return nil, err
-        end
-        conn = (conn_from_red or 0) + 1
+        red:zremrangebyscore(key, 0, ngx_time())
+        local count, err = red:zcard(key)
+        if err then return nil, err end
+        conn = (count or 0) + 1
     end
 
     if conn > max then
@@ -60,11 +124,19 @@ function _M.incoming(self, red, key, commit)
 end
 
 
-function _M.leaving(self, red, key, req_latency)
+function _M.leaving(self, red, key, req_latency, req_id)
     assert(key)
     key = "limit_conn" .. ":" .. key
 
-    local conn, err = red:incrby(key, -1)
+    local conn, err
+    if req_id then
+        local res, err = red:zrem(key, req_id)
+        if not res then
+            return nil, err
+        end
+    end
+    conn, err = red:zcard(key)
+
     if not conn then
         return nil, err
     end

apisix/utils/redis-schema.lua

@@ -74,8 +74,20 @@ local policy_to_additional_properties = {
     },
 }
 
+local limit_conn_redis_cluster_schema = policy_to_additional_properties["redis-cluster"]
+limit_conn_redis_cluster_schema.properties.key_ttl = {
+    type = "integer", default = 3600,
+}
+
+local limit_conn_redis_schema = policy_to_additional_properties["redis"]
+limit_conn_redis_schema.properties.key_ttl = {
+    type = "integer", default = 3600,
+}
+
 local _M = {
-    schema = policy_to_additional_properties
+    schema = policy_to_additional_properties,
+    limit_conn_redis_cluster_schema = limit_conn_redis_cluster_schema,
+    limit_conn_redis_schema = limit_conn_redis_schema,
 }
 
 return _M

docs/en/latest/plugins/limit-conn.md

@@ -44,6 +44,7 @@ The `limit-conn` Plugin limits the rate of requests by the number of concurrent
 | only_use_default_delay   | boolean | False    | false       |      | If false, delay requests proportionally based on how much they exceed the `conn` limit. The delay grows larger as congestion increases. For instance, with `conn` being `5`, `burst` being `3`, and `default_conn_delay` being `1`, 6 concurrent requests would result in a 1-second delay, 7 requests a 2-second delay, 8 requests a 3-second delay, and so on, until the total limit of `conn + burst` is reached, beyond which requests are rejected. If true, use `default_conn_delay` to delay all excessive requests within the `burst` range. Requests beyond `conn + burst` are rejected immediately. For instance, with `conn` being `5`, `burst` being `3`, and `default_conn_delay` being `1`, 6, 7, or 8 concurrent requests are all delayed by exactly 1 second each. |
 | key_type        | string  | False      | var   | ["var","var_combination"] | The type of key. If the `key_type` is `var`, the `key` is interpreted a variable. If the `key_type` is `var_combination`, the `key` is interpreted as a combination of variables.    |
 | key       | string  | False      | remote_addr |   | The key to count requests by. If the `key_type` is `var`, the `key` is interpreted a variable. The variable does not need to be prefixed by a dollar sign (`$`). If the `key_type` is `var_combination`, the `key` is interpreted as a combination of variables. All variables should be prefixed by dollar signs (`$`). For example, to configure the `key` to use a combination of two request headers `custom-a` and `custom-b`, the `key` should be configured as `$http_custom_a $http_custom_b`. |
+| key_ttl   | integer | False      | 3600          |   | The TTL of the Redis key in seconds. Used when `policy` is `redis` or `redis-cluster`. |
 | rejected_code   | integer | False      | 503   | [200,...,599]   | The HTTP status code returned when a request is rejected for exceeding the threshold.     |
 | rejected_msg    | string  | False        |       | non-empty   | The response body returned when a request is rejected for exceeding the threshold.     |
 | allow_degradation       | boolean | False      | false   |   | If true, allow APISIX to continue handling requests without the Plugin when the Plugin or its dependencies become unavailable.        |

docs/zh/latest/plugins/limit-conn.md

@@ -44,6 +44,7 @@ description: limit-conn 插件通过管理并发连接来限制请求速率。
 | only_use_default_delay | boolean | 否 | false | | 如果为 false，则根据请求超出`conn`限制的程度按比例延迟请求。拥塞越严重，延迟就越大。例如，当 `conn` 为 `5`、`burst` 为 `3` 且 `default_conn_delay` 为 `1` 时，6 个并发请求将导致 1 秒的延迟，7 个请求将导致 2 秒的延迟，8 个请求将导致 3 秒的延迟，依此类推，直到达到 `conn + burst` 的总限制，超过此限制的请求将被拒绝。如果为 true，则使用 `default_conn_delay` 延迟 `burst` 范围内的所有超额请求。超出 `conn + burst` 的请求将被立即拒绝。例如，当 `conn` 为 `5`、`burst` 为 `3` 且 `default_conn_delay` 为 `1` 时，6、7 或 8 个并发请求都将延迟 1 秒。|
 | key_type | string | 否 | var | ["var","var_combination"] | key 的类型。如果`key_type` 为 `var`，则 `key` 将被解释为变量。如果 `key_type` 为 `var_combination`，则 `key` 将被解释为变量的组合。 |
 | key | string | 否 | remote_addr | | 用于计数请求的 key。如果 `key_type` 为 `var`，则 `key` 将被解释为变量。变量不需要以美元符号（`$`）为前缀。如果 `key_type` 为 `var_combination`，则 `key` 会被解释为变量的组合。所有变量都应该以美元符号 (`$`) 为前缀。例如，要配置 `key` 使用两个请求头 `custom-a` 和 `custom-b` 的组合，则 `key` 应该配置为 `$http_custom_a $http_custom_b`。|
+| key_ttl | integer | 否 | 3600 | | Redis 键的 TTL（以秒为单位）。当 `policy` 为 `redis` 或 `redis-cluster` 时使用。 |
 | rejection_code | integer | 否 | 503 | [200,...,599] | 请求因超出阈值而被拒绝时返回的 HTTP 状态代码。|
 | rejection_msg | string | 否 | | 非空 | 请求因超出阈值而被拒绝时返回的响应主体。|
 | allow_degradation | boolean | 否 | false | | 如果为 true，则允许 APISIX 在插件或其依赖项不可用时继续处理没有插件的请求。|