feat: add redis and redis-cluster in limit-conn (#10866)

Author: theweakgod

apisix/cli/ngx_tpl.lua

@@ -289,6 +289,7 @@ http {
 
     {% if enabled_plugins["limit-conn"] then %}
     lua_shared_dict plugin-limit-conn {* http.lua_shared_dict["plugin-limit-conn"] *};
+    lua_shared_dict plugin-limit-conn-redis-cluster-slot-lock {* http.lua_shared_dict["plugin-limit-conn-redis-cluster-slot-lock"] *};
     {% end %}
 
     {% if enabled_plugins["limit-req"] then %}

apisix/plugins/limit-conn.lua

@@ -14,15 +14,18 @@
 -- See the License for the specific language governing permissions and
 -- limitations under the License.
 --
-local core = require("apisix.core")
-local limit_conn = require("apisix.plugins.limit-conn.init")
+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 plugin_name = "limit-conn"
 local schema = {
     type = "object",
     properties = {
-        conn = {type = "integer", exclusiveMinimum = 0},
+        conn = {type = "integer", exclusiveMinimum = 0},               -- limit.conn max
         burst = {type = "integer",  minimum = 0},
         default_conn_delay = {type = "number", exclusiveMinimum = 0},
         only_use_default_delay = {type = "boolean", default = false},
@@ -31,6 +34,11 @@ local schema = {
             enum = {"var", "var_combination"},
             default = "var",
         },
+        policy = {
+            type = "string",
+            enum = {"redis", "redis-cluster", "local"},
+            default = "local",
+        },
         rejected_code = {
             type = "integer", minimum = 200, maximum = 599, default = 503
         },
@@ -39,7 +47,25 @@ local schema = {
         },
         allow_degradation = {type = "boolean", default = false}
     },
-    required = {"conn", "burst", "default_conn_delay", "key"}
+    required = {"conn", "burst", "default_conn_delay", "key"},
+    ["if"] = {
+        properties = {
+            policy = {
+                enum = {"redis"},
+            },
+        },
+    },
+    ["then"] = policy_to_additional_properties.redis,
+    ["else"] = {
+        ["if"] = {
+            properties = {
+                policy = {
+                    enum = {"redis-cluster"},
+                },
+            },
+        },
+        ["then"] = policy_to_additional_properties["redis-cluster"],
+    }
 }
 
 local _M = {

apisix/plugins/limit-conn/init.lua

@@ -23,6 +23,16 @@ if ngx.config.subsystem == "stream" then
     shdict_name = shdict_name .. "-stream"
 end
 
+local redis_single_new
+local redis_cluster_new
+do
+    local redis_src = "apisix.plugins.limit-conn.limit-conn-redis"
+    redis_single_new = require(redis_src).new
+
+    local cluster_src = "apisix.plugins.limit-conn.limit-conn-redis-cluster"
+    redis_cluster_new = require(cluster_src).new
+end
+
 
 local lrucache = core.lrucache.new({
     type = "plugin",
@@ -31,9 +41,26 @@ local _M = {}
 
 
 local function create_limit_obj(conf)
-    core.log.info("create new limit-conn plugin instance")
-    return limit_conn_new(shdict_name, conf.conn, conf.burst,
-                          conf.default_conn_delay)
+    if conf.policy == "local" then
+        core.log.info("create new limit-conn plugin instance")
+        return limit_conn_new(shdict_name, conf.conn, conf.burst,
+                              conf.default_conn_delay)
+    elseif 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,
+                                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,
+                                 conf.default_conn_delay)
+    else
+        return nil, "policy enum not match"
+    end
 end
 
 

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

@@ -0,0 +1,78 @@
+--
+-- 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.
+--
+local redis_cluster     = require("apisix.utils.rediscluster")
+local core              = require("apisix.core")
+local util              = require("apisix.plugins.limit-conn.util")
+local setmetatable      = setmetatable
+local ngx_timer_at      = ngx.timer.at
+
+local _M = {version = 0.1}
+
+
+local mt = {
+    __index = _M
+}
+
+
+function _M.new(plugin_name, conf, max, burst, default_conn_delay)
+
+    local red_cli, err = redis_cluster.new(conf, "plugin-limit-conn-redis-cluster-slot-lock")
+    if not red_cli then
+        return nil, err
+    end
+    local self = {
+        conf = conf,
+        plugin_name = plugin_name,
+        burst = burst,
+        max = max + 0,    -- just to ensure the param is good
+        unit_delay = default_conn_delay,
+        red_cli = red_cli,
+    }
+    return setmetatable(self, mt)
+end
+
+
+function _M.incoming(self, key, commit)
+    return util.incoming(self, self.red_cli, key, commit)
+end
+
+
+function _M.is_committed(self)
+    return self.committed
+end
+
+
+local function leaving_thread(premature, self, key, req_latency)
+    return util.leaving(self, self.red_cli, key, req_latency)
+end
+
+
+function _M.leaving(self, key, req_latency)
+    -- log_by_lua can't use cosocket
+    local ok, err = ngx_timer_at(0, leaving_thread, self, key, req_latency)
+    if not ok then
+        core.log.error("failed to create timer: ", err)
+        return nil, err
+    end
+
+    return ok
+
+end
+
+
+
+return _M

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

@@ -0,0 +1,85 @@
+--
+-- 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.
+--
+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 setmetatable  = setmetatable
+
+
+local _M = {version = 0.1}
+
+
+local mt = {
+    __index = _M
+}
+
+function _M.new(plugin_name, conf, max, burst, default_conn_delay)
+
+    local self = {
+        conf = conf,
+        plugin_name = plugin_name,
+        burst = burst,
+        max = max + 0,    -- just to ensure the param is good
+        unit_delay = default_conn_delay,
+    }
+    return setmetatable(self, mt)
+end
+
+
+function _M.incoming(self, key, commit)
+    local conf = self.conf
+    local red, err = redis.new(conf)
+    if not red then
+        return red, err
+    end
+    return util.incoming(self, red, key, commit)
+end
+
+
+function _M.is_committed(self)
+    return self.committed
+end
+
+
+local function leaving_thread(premature, self, key, req_latency)
+
+    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)
+end
+
+
+function _M.leaving(self, key, req_latency)
+    -- log_by_lua can't use cosocket
+    local ok, err = ngx_timer_at(0, leaving_thread, self, key, req_latency)
+    if not ok then
+        core.log.error("failed to create timer: ", err)
+        return nil, err
+    end
+
+    return ok
+
+end
+
+
+
+return _M

apisix/plugins/limit-conn/util.lua

@@ -0,0 +1,81 @@
+--
+-- 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.
+--
+
+local assert            = assert
+local math              = require "math"
+local floor             = math.floor
+local _M = {version = 0.3}
+
+
+function _M.incoming(self, red, key, commit)
+    local max = self.max
+    self.committed = false
+    key = "limit_conn" .. ":" .. key
+
+    local conn, err
+    if commit then
+        conn, err = red:incrby(key, 1)
+        if not conn then
+            return nil, err
+        end
+
+        if conn > max + self.burst then
+            conn, err = red:incrby(key, -1)
+            if not conn then
+                return nil, err
+            end
+            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
+    end
+
+    if conn > max then
+        -- make the excessive connections wait
+        return self.unit_delay * floor((conn - 1) / max), conn
+    end
+
+    -- we return a 0 delay by default
+    return 0, conn
+end
+
+
+function _M.leaving(self, red, key, req_latency)
+    assert(key)
+    key = "limit_conn" .. ":" .. key
+
+    local conn, err = red:incrby(key, -1)
+    if not conn then
+        return nil, err
+    end
+
+    if req_latency then
+        local unit_delay = self.unit_delay
+        self.unit_delay = (req_latency + unit_delay) / 2
+    end
+
+    return conn
+end
+
+
+return _M