feat(stream): add limit-conn

Makefile

@@ -215,6 +215,9 @@ install: default
 	$(INSTALL) -d $(INST_LUADIR)/apisix/plugins/grpc-transcode
 	$(INSTALL) apisix/plugins/grpc-transcode/*.lua $(INST_LUADIR)/apisix/plugins/grpc-transcode/
 
+	$(INSTALL) -d $(INST_LUADIR)/apisix/plugins/limit-conn
+	$(INSTALL) apisix/plugins/limit-conn/*.lua $(INST_LUADIR)/apisix/plugins/limit-conn/
+
 	$(INSTALL) -d $(INST_LUADIR)/apisix/plugins/limit-count
 	$(INSTALL) apisix/plugins/limit-count/*.lua $(INST_LUADIR)/apisix/plugins/limit-count/
 

apisix/cli/ngx_tpl.lua

@@ -67,6 +67,7 @@ stream {
     lua_socket_log_errors off;
 
     lua_shared_dict lrucache-lock-stream {* stream.lua_shared_dict["lrucache-lock-stream"] *};
+    lua_shared_dict plugin-limit-conn-stream {* stream.lua_shared_dict["plugin-limit-conn-stream"] *};
 
     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 = shdict_name .. "-stream"
+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

@@ -148,6 +148,7 @@ nginx_config:                     # config for render the template to generate n
   stream:
     lua_shared_dict:
       lrucache-lock-stream: 10m
+      plugin-limit-conn-stream: 10m
 
   # As user can add arbitrary configurations in the snippet,
   # it is user's responsibility to check the configurations
@@ -319,6 +320,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,8 @@ title: limit-conn
 
 **注：key 是可以被用户自定义的，只需要修改插件的一行代码即可完成。并没有在插件中放开是处于安全的考虑。**
 
+在 stream 代理中使用该插件时，只有 `remote_addr` 和 `server_addr` 可以被用作 key。另外设置 `rejected_code` 毫无意义。
+
 #### 如何启用
 
 下面是一个示例，在指定的 route 上开启了 limit-conn 插件:

t/APISIX.pm

@@ -295,11 +295,18 @@ _EOC_
     my $stream_enable = $block->stream_enable;
     my $stream_conf_enable = $block->stream_conf_enable;
     my $extra_stream_config = $block->extra_stream_config // '';
+    my $stream_upstream_code = $block->stream_upstream_code // <<_EOC_;
+            local sock = ngx.req.socket()
+            local data = sock:receive("1")
+            ngx.say("hello world")
+_EOC_
+
     my $stream_config = $block->stream_config // <<_EOC_;
     $lua_deps_path
     lua_socket_log_errors off;
 
     lua_shared_dict lrucache-lock-stream 10m;
+    lua_shared_dict plugin-limit-conn-stream 10m;
 
     upstream apisix_backend {
         server 127.0.0.1:1900;
@@ -339,9 +346,7 @@ _EOC_
         listen 1995;
 
         content_by_lua_block {
-            local sock = ngx.req.socket()
-            local data = sock:receive("1")
-            ngx.say("hello world")
+            $stream_upstream_code
         }
     }
 _EOC_
@@ -536,20 +541,6 @@ _EOC_
                 apisix.http_log_phase()
             }
         }
-
-        location = /v3/auth/authenticate {
-            content_by_lua_block {
-                ngx.log(ngx.WARN, "etcd auth failed!")
-            }
-        }
-
-        location  = /.well-known/openid-configuration {
-            content_by_lua_block {
-                ngx.say([[
-{"issuer":"https://samples.auth0.com/","authorization_endpoint":"https://samples.auth0.com/authorize","token_endpoint":"https://samples.auth0.com/oauth/token","device_authorization_endpoint":"https://samples.auth0.com/oauth/device/code","userinfo_endpoint":"https://samples.auth0.com/userinfo","mfa_challenge_endpoint":"https://samples.auth0.com/mfa/challenge","jwks_uri":"https://samples.auth0.com/.well-known/jwks.json","registration_endpoint":"https://samples.auth0.com/oidc/register","revocation_endpoint":"https://samples.auth0.com/oauth/revoke","scopes_supported":["openid","profile","offline_access","name","given_name","family_name","nickname","email","email_verified","picture","created_at","identities","phone","address"],"response_types_supported":["code","token","id_token","code token","code id_token","token id_token","code token id_token"],"code_challenge_methods_supported":["S256","plain"],"response_modes_supported":["query","fragment","form_post"],"subject_types_supported":["public"],"id_token_signing_alg_values_supported":["HS256","RS256"],"token_endpoint_auth_methods_supported":["client_secret_basic","client_secret_post"],"claims_supported":["aud","auth_time","created_at","email","email_verified","exp","family_name","given_name","iat","identities","iss","name","nickname","phone_number","picture","sub"],"request_uri_parameter_supported":false}
-                ]])
-            }
-        }
     }
 
     $a6_ngx_directives