feat(proxy-mirror): support mirror requests sample_ratio

apisix/cli/ngx_tpl.lua

@@ -188,6 +188,9 @@ http {
     # for authz-keycloak
     lua_shared_dict access-tokens {* http.lua_shared_dict["access-tokens"] *}; # cache for service account access tokens
 
+    # for proxy-mirror plugin
+    lua_shared_dict plugin-proxy-mirror {* http.lua_shared_dict["plugin-proxy-mirror"] *};
+
     # for custom shared dict
     {% if http.lua_shared_dicts then %}
     {% for cache_key, cache_size in pairs(http.lua_shared_dicts) do %}

apisix/plugins/proxy-mirror.lua

@@ -15,8 +15,20 @@
 -- limitations under the License.
 --
 local core          = require("apisix.core")
-
+local limit_local_new = require("resty.limit.count").new
 local plugin_name   = "proxy-mirror"
+local limit_redis_cluster_new
+local limit_redis_new
+do
+    local redis_src = "apisix.plugins.limit-count.limit-count-redis"
+    limit_redis_new = require(redis_src).new
+
+    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 schema = {
     type = "object",
@@ -26,9 +38,77 @@ local schema = {
             pattern = [[^http(s)?:\/\/[a-zA-Z0-9][-a-zA-Z0-9]{0,62}]]
                       .. [[(\.[a-zA-Z0-9][-a-zA-Z0-9]{0,62})+(:[0-9]{1,5})?$]],
         },
+        count = { type = "integer", minimum = 1, default = 1 },
+        time_window = { type = "integer", minimum = 1, default = 60 },
+        policy = {
+            type = "string",
+            enum = { "local", "redis", "redis-cluster" },
+            default = "local",
+        },
+        enable_limit_count = { type = "boolean", default = false },
     },
     required = {"host"},
     minProperties = 1,
+    dependencies = {
+        policy = {
+            oneOf = {
+                {
+                    properties = {
+                        policy = {
+                            enum = {"local"},
+                        },
+                    },
+                },
+                {
+                    properties = {
+                        policy = {
+                            enum = {"redis"},
+                        },
+                        redis_host = {
+                            type = "string", minLength = 2,
+                        },
+                        redis_port = {
+                            type = "integer", minimum = 1, default = 6379,
+                        },
+                        redis_password = {
+                            type = "string", minLength = 0,
+                        },
+                        redis_database = {
+                            type = "integer", minimum = 0, default = 0,
+                        },
+                        redis_timeout = {
+                            type = "integer", minimum = 1, default = 1000,
+                        },
+                    },
+                    required = {"redis_host"},
+                },
+                {
+                    properties = {
+                        policy = {
+                            enum = {"redis-cluster"},
+                        },
+                        redis_cluster_nodes = {
+                            type = "array",
+                            minItems = 2,
+                            items = {
+                                type = "string", minLength = 2, maxLength = 100,
+                            },
+                        },
+                        redis_password = {
+                            type = "string", minLength = 0,
+                        },
+                        redis_timeout = {
+                            type = "integer", minimum = 1, default = 1000,
+                        },
+                        redis_cluster_name = {
+                            type = "string",
+                        },
+                    },
+                    required = {"redis_cluster_nodes", "redis_cluster_name"},
+                }
+            }
+        }
+    }
 }
 
 local _M = {
@@ -48,12 +128,57 @@ function _M.check_schema(conf)
     return true
 end
 
+local function create_limit_obj(conf)
+    core.log.info("create new limit count object instance")
+
+    if not conf.policy or conf.policy == "local" then
+        return limit_local_new("plugin-" .. plugin_name, conf.count,
+            conf.time_window)
+    end
+
+    if conf.policy == "redis" then
+        return limit_redis_new("plugin-" .. plugin_name,
+            conf.count, conf.time_window, conf)
+    end
+
+    if conf.policy == "redis-cluster" then
+        return limit_redis_cluster_new("plugin-" .. plugin_name, conf.count,
+            conf.time_window, conf)
+    end
+
+    return nil
+end
 
 function _M.rewrite(conf, ctx)
     core.log.info("proxy mirror plugin rewrite phase, conf: ", core.json.delay_encode(conf))
 
     ctx.var.upstream_host = ctx.var.host
-    ctx.var.upstream_mirror_host = conf.host
+
+    local is_mirror_request = true
+    if conf.enable_limit_count then
+        local lim, err = core.lrucache.plugin_ctx(lrucache, ctx, conf.policy,
+            create_limit_obj, conf)
+        if not lim then
+            core.log.error("failed to fetch limit count object: ", err)
+        else
+            local key = ctx.conf_type .. ctx.conf_version
+            core.log.info("limit key: ", key)
+
+            local delay, remaining = lim:incoming(key, true)
+            if not delay then
+                local err = remaining
+                if err == "rejected" then
+                    is_mirror_request = false
+                else
+                    core.log.error("failed to limit mirror request: ", err)
+                end
+            end
+        end
+    end
+
+    if is_mirror_request then
+        ctx.var.upstream_mirror_host = conf.host
+    end
 end
 
 

conf/config-default.yaml

@@ -235,6 +235,7 @@ nginx_config:                     # config for render the template to generate n
       jwks: 1m
       introspection: 10m
       access-tokens: 1m
+      plugin-proxy-mirror: 10m
 
 etcd:
   host:                           # it's possible to define multiple etcd hosts addresses of the same etcd cluster.

docs/en/latest/plugins/proxy-mirror.md

@@ -29,13 +29,24 @@ The proxy-mirror plugin, which provides the ability to mirror client requests.
 
 | Name | Type   | Requirement | Default | Valid | Description                                                                                                                 |
 | ---- | ------ | ----------- | ------- | ----- | --------------------------------------------------------------------------------------------------------------------------- |
-| host | string | optional    |         |       | Specify a mirror service address, e.g. http://127.0.0.1:9797 (address needs to contain schema: http or https, not URI part) |
+| host | string | required    |         |       | Specify a mirror service address, e.g. http://127.0.0.1:9797 (address needs to contain schema: http or https, not URI part) |
+| enable_limit_count               | boolean | optional                                | false              |                                                                                                | Whether to enable the limit count of mirror requests.                                                                                                                                                                                                                                                                |
+| count               | integer | optional                                | 1              | count > 0                                                                                               | the specified number of mirror requests threshold.                                                                                                                                                                                                                                                                |
+| time_window         | integer | optional                    | 60              | time_window > 0                                                                                         | the time window in seconds before the mirror request count is reset.                                                                                                                                                                                                                                              |
+| policy              | string  | optional                                | "local"       | ["local", "redis", "redis-cluster"]                                                                     | The mirror request limit count policies to use for retrieving and incrementing the limits. Available values are `local`(the counters will be stored locally in-memory on the node), `redis`(counters are stored on a Redis server and will be shared across the nodes, usually use it to do the global speed limit), and `redis-cluster` which works the same as `redis` but with redis cluster. |
+| redis_host          | string  | required for `redis`                    |               |                                                                                                         | When using the `redis` policy, this property specifies the address of the Redis server.                                                                                                                                                                                                                    |
+| redis_port          | integer | optional                                | 6379          | [1,...]                                                                                                 | When using the `redis` policy, this property specifies the port of the Redis server.                                                                                                                                                                                                                       |
+| redis_password      | string  | optional                                |               |                                                                                                         | When using the `redis`  or `redis-cluster` policy, this property specifies the password of the Redis server.                                                                                                                                                                                                                   |
+| redis_database      | integer | optional                                | 0             | redis_database >= 0                                                                                     | When using the `redis` policy, this property specifies the database you selected of the Redis server, and only for non Redis cluster mode (single instance mode or Redis public cloud service that provides single entry).                                                                               |
+| redis_timeout       | integer | optional                                | 1000          | [1,...]                                                                                                 | When using the `redis`  or `redis-cluster` policy, this property specifies the timeout in milliseconds of any command submitted to the Redis server.                                                                                                                                                                           |
+| redis_cluster_nodes | array   | required when policy is `redis-cluster` |               |                                                                                                         | When using `redis-cluster` policy，This property is a list of addresses of Redis cluster service nodes (at least two).                                                                                                                                                                                                    |
+| redis_cluster_name  | string  | required when policy is `redis-cluster` |               |                                                                                                         | When using `redis-cluster` policy, this property is the name of Redis cluster service nodes.                                                                                                                                                                                                                   |
 
 ### Examples
 
 #### Enable the plugin
 
-1:  enable the proxy-mirror plugin for a specific route :
+example 1:  enable the proxy-mirror plugin for a specific route:
 
 ```shell
 curl http://127.0.0.1:9080/apisix/admin/routes/1  -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
@@ -73,6 +84,31 @@ hello world
 > Since the specified mirror address is 127.0.0.1:9797, so to verify whether this plugin is in effect, we need to confirm on the service with port 9797.
 > For example, we can start a simple server:  python -m SimpleHTTPServer 9797
 
+example 2:  enable the proxy-mirror plugin for a specific route and enable limit count for mirror requests:
+
+```shell
+curl http://127.0.0.1:9080/apisix/admin/routes/1  -H 'X-API-KEY: edd1c9f034335f136f87ad84b625c8f1' -X PUT -d '
+{
+    "plugins": {
+        "proxy-mirror": {
+           "host": "http://127.0.0.1:9797",
+            "enable_limit_count": true,
+            "count": 2,
+            "time_window": 60
+        }
+    },
+    "upstream": {
+        "nodes": {
+            "127.0.0.1:1999": 1
+        },
+        "type": "roundrobin"
+    },
+    "uri": "/hello"
+}'
+```
+
+In this case, only allow two mirror requests in 60 seconds.
+
 ## Disable Plugin
 
 Remove the corresponding JSON in the plugin configuration to disable the plugin immediately without restarting the service: