feat(stream): add limit-conn

Signed-off-by: spacewander <[email protected]>

apisix/cli/ngx_tpl.lua

@@ -67,6 +67,7 @@ stream {
     lua_socket_log_errors off;
 
     lua_shared_dict lrucache-lock-stream   10m;
+    lua_shared_dict stream-plugin-limit-conn 10m;
 
     resolver {% for _, dns_addr in ipairs(dns_resolver or {}) do %} {*dns_addr*} {% end %} {% if dns_resolver_valid then %} valid={*dns_resolver_valid*}{% end %};
     resolver_timeout {*resolver_timeout*};

apisix/plugins/limit-conn.lua

@@ -14,16 +14,11 @@
 -- See the License for the specific language governing permissions and
 -- limitations under the License.
 --
-local limit_conn_new = require("resty.limit.conn").new
 local core = require("apisix.core")
-local sleep = core.sleep
-local plugin_name = "limit-conn"
-
+local limit_conn = require("apisix.plugins.limit-conn.init")
 
-local lrucache = core.lrucache.new({
-    type = "plugin",
-})
 
+local plugin_name = "limit-conn"
 local schema = {
     type = "object",
     properties = {
@@ -41,95 +36,26 @@ local schema = {
     required = {"conn", "burst", "default_conn_delay", "key"}
 }
 
-
 local _M = {
     version = 0.1,
     priority = 1003,
     name = plugin_name,
     schema = schema,
 }
 
-function _M.check_schema(conf)
-    local ok, err = core.schema.check(schema, conf)
-    if not ok then
-        return false, err
-    end
-
-    return true
-end
 
-local function create_limit_obj(conf)
-    core.log.info("create new limit-conn plugin instance")
-    return limit_conn_new("plugin-limit-conn", conf.conn, conf.burst,
-                          conf.default_conn_delay)
+function _M.check_schema(conf)
+    return core.schema.check(schema, conf)
 end
 
 
 function _M.access(conf, ctx)
-    core.log.info("ver: ", ctx.conf_version)
-    local lim, err = lrucache(conf, nil, create_limit_obj, conf)
-    if not lim then
-        core.log.error("failed to instantiate a resty.limit.conn object: ", err)
-        return 500
-    end
-
-    local key = (ctx.var[conf.key] or "") .. ctx.conf_type .. ctx.conf_version
-    core.log.info("limit key: ", key)
-
-    local delay, err = lim:incoming(key, true)
-    if not delay then
-        if err == "rejected" then
-            return conf.rejected_code
-        end
-
-        core.log.error("failed to limit req: ", err)
-        return 500
-    end
-
-    if lim:is_committed() then
-        if not ctx.limit_conn then
-            ctx.limit_conn = core.tablepool.fetch("plugin#limit-conn", 0, 6)
-        end
-
-        core.table.insert_tail(ctx.limit_conn, lim, key, delay)
-    end
-
-    if delay >= 0.001 then
-        sleep(delay)
-    end
+    return limit_conn.increase(conf, ctx)
 end
 
 
 function _M.log(conf, ctx)
-    local limit_conn = ctx.limit_conn
-    if not limit_conn then
-        return
-    end
-
-    for i = 1, #limit_conn, 3 do
-        local lim = limit_conn[i]
-        local key = limit_conn[i + 1]
-        local delay = limit_conn[i + 2]
-
-        local latency
-        if ctx.proxy_passed then
-            latency = ctx.var.upstream_response_time
-        else
-            latency = ctx.var.request_time - delay
-        end
-
-        core.log.debug("request latency is ", latency) -- for test
-
-        local conn, err = lim:leaving(key, latency)
-        if not conn then
-            core.log.error("failed to record the connection leaving request: ",
-                           err)
-            break
-        end
-    end
-
-    core.tablepool.release("plugin#limit-conn", limit_conn)
-    return
+    return limit_conn.decrease(conf, ctx)
 end
 
 

apisix/plugins/limit-conn/init.lua

@@ -0,0 +1,107 @@
+--
+-- 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 limit_conn_new = require("resty.limit.conn").new
+local core = require("apisix.core")
+local sleep = core.sleep
+local shdict_name = "plugin-limit-conn"
+if ngx.config.subsystem == "stream" then
+    shdict_name = "stream-" .. shdict_name
+end
+
+
+local lrucache = core.lrucache.new({
+    type = "plugin",
+})
+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)
+end
+
+
+function _M.increase(conf, ctx)
+    core.log.info("ver: ", ctx.conf_version)
+    local lim, err = lrucache(conf, nil, create_limit_obj, conf)
+    if not lim then
+        core.log.error("failed to instantiate a resty.limit.conn object: ", err)
+        return 500
+    end
+
+    local key = (ctx.var[conf.key] or "") .. ctx.conf_type .. ctx.conf_version
+    core.log.info("limit key: ", key)
+
+    local delay, err = lim:incoming(key, true)
+    if not delay then
+        if err == "rejected" then
+            return conf.rejected_code or 503
+        end
+
+        core.log.error("failed to limit req: ", err)
+        return 500
+    end
+
+    if lim:is_committed() then
+        if not ctx.limit_conn then
+            ctx.limit_conn = core.tablepool.fetch("plugin#limit-conn", 0, 6)
+        end
+
+        core.table.insert_tail(ctx.limit_conn, lim, key, delay)
+    end
+
+    if delay >= 0.001 then
+        sleep(delay)
+    end
+end
+
+
+function _M.decrease(conf, ctx)
+    local limit_conn = ctx.limit_conn
+    if not limit_conn then
+        return
+    end
+
+    for i = 1, #limit_conn, 3 do
+        local lim = limit_conn[i]
+        local key = limit_conn[i + 1]
+        local delay = limit_conn[i + 2]
+
+        local latency
+        if ctx.proxy_passed then
+            latency = ctx.var.upstream_response_time
+        else
+            latency = ctx.var.request_time - delay
+        end
+
+        core.log.debug("request latency is ", latency) -- for test
+
+        local conn, err = lim:leaving(key, latency)
+        if not conn then
+            core.log.error("failed to record the connection leaving request: ",
+                           err)
+            break
+        end
+    end
+
+    core.tablepool.release("plugin#limit-conn", limit_conn)
+    return
+end
+
+
+return _M

apisix/stream/plugins/limit-conn.lua

@@ -0,0 +1,59 @@
+--
+-- 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 core = require("apisix.core")
+local limit_conn = require("apisix.plugins.limit-conn.init")
+
+
+local plugin_name = "limit-conn"
+local schema = {
+    type = "object",
+    properties = {
+        conn = {type = "integer", exclusiveMinimum = 0},
+        burst = {type = "integer",  minimum = 0},
+        default_conn_delay = {type = "number", exclusiveMinimum = 0},
+        key = {
+            type = "string",
+            enum = {"remote_addr", "server_addr"}
+        },
+    },
+    required = {"conn", "burst", "default_conn_delay", "key"}
+}
+
+local _M = {
+    version = 0.1,
+    priority = 1003,
+    name = plugin_name,
+    schema = schema,
+}
+
+
+function _M.check_schema(conf)
+    return core.schema.check(schema, conf)
+end
+
+
+function _M.preread(conf, ctx)
+    return limit_conn.increase(conf, ctx)
+end
+
+
+function _M.log(conf, ctx)
+    return limit_conn.decrease(conf, ctx)
+end
+
+
+return _M

conf/config-default.yaml

@@ -293,6 +293,7 @@ plugins:                          # plugin list (sorted by priority)
   - ext-plugin-post-req            # priority: -3000
 
 stream_plugins: # sorted by priority
+  - limit-conn                     # priority: 1003
   - mqtt-proxy                     # priority: 1000
   # <- recommend to use priority (0, 100) for your custom plugins
 

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

@@ -45,6 +45,8 @@ Limiting request concurrency plugin.
 
 **Key can be customized by the user, only need to modify a line of code of the plug-in to complete. It is a security consideration that is not open in the plugin.**
 
+When used in the stream proxy, only `remote_addr` and `server_addr` can be used as key. And `rejected_code` is meaningless.
+
 ## How To Enable
 
 Here's an example, enable the limit-conn plugin on the specified route:

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

@@ -35,6 +35,9 @@ title: limit-conn
 
 **注：key 是可以被用户自定义的，只需要修改插件的一行代码即可完成。并没有在插件中放开是处于安全的考虑。**
 
+: when used in the stream proxy, only `remote_addr` and `server_addr` can be used as `key`. And `rejected_code` is meaningless.
+在 stream 代理中使用该插件时，只有 `remote_addr` 和 `server_addr` 可以被用作 key。另外设置 `rejected_code` 毫无意义。
+
 #### 如何启用
 
 下面是一个示例，在指定的 route 上开启了 limit-conn 插件:

t/APISIX.pm

@@ -299,6 +299,7 @@ _EOC_
     lua_socket_log_errors off;
 
     lua_shared_dict lrucache-lock-stream   10m;
+    lua_shared_dict stream-plugin-limit-conn 10m;
 
     upstream apisix_backend {
         server 127.0.0.1:1900;

t/plugin/limit-conn2.t

@@ -16,7 +16,7 @@
 #
 BEGIN {
     if ($ENV{TEST_NGINX_CHECK_LEAK}) {
-        $SkipReason = "unavailable for the hup tests";
+        $SkipReason = "unavailable for the check leak tests";
 
     } else {
         $ENV{TEST_NGINX_USE_HUP} = 1;