feat: limit count plugin support X-RateLimit-Reset

apisix/cli/ngx_tpl.lua

@@ -276,6 +276,7 @@ http {
     {% if enabled_plugins["limit-count"] then %}
     lua_shared_dict plugin-limit-count {* http.lua_shared_dict["plugin-limit-count"] *};
     lua_shared_dict plugin-limit-count-redis-cluster-slot-lock {* http.lua_shared_dict["plugin-limit-count-redis-cluster-slot-lock"] *};
+    lua_shared_dict plugin-limit-count-reset-header {* http.lua_shared_dict["plugin-limit-count"] *};
     {% end %}
 
     {% if enabled_plugins["prometheus"] and not enabled_stream_plugins["prometheus"] then %}

apisix/plugins/limit-count/init.lua

@@ -14,18 +14,20 @@
 -- See the License for the specific language governing permissions and
 -- limitations under the License.
 --
-local limit_local_new = require("resty.limit.count").new
 local core = require("apisix.core")
 local apisix_plugin = require("apisix.plugin")
 local tab_insert = table.insert
 local ipairs = ipairs
 local pairs = pairs
 
-
 local plugin_name = "limit-count"
 local limit_redis_cluster_new
 local limit_redis_new
+local limit_local_new
 do
+    local local_src = "apisix.plugins.limit-count.limit-count-local"
+    limit_local_new = require(local_src).new
+
     local redis_src = "apisix.plugins.limit-count.limit-count-redis"
     limit_redis_new = require(redis_src).new
 
@@ -39,7 +41,6 @@ local group_conf_lru = core.lrucache.new({
     type = 'plugin',
 })
 
-
 local policy_to_additional_properties = {
     redis = {
         properties = {
@@ -242,7 +243,6 @@ local function gen_limit_obj(conf, ctx)
     return core.lrucache.plugin_ctx(lrucache, ctx, extra_key, create_limit_obj, conf)
 end
 
-
 function _M.rate_limit(conf, ctx)
     core.log.info("ver: ", ctx.conf_version)
 
@@ -283,10 +283,17 @@ function _M.rate_limit(conf, ctx)
     key = gen_limit_key(conf, ctx, key)
     core.log.info("limit key: ", key)
 
-    local delay, remaining = lim:incoming(key, true)
+    local delay, remaining, reset = lim:incoming(key, true, conf)
     if not delay then
         local err = remaining
         if err == "rejected" then
+            -- show count limit header when rejected
+            if conf.show_limit_quota_header then
+                core.response.set_header("X-RateLimit-Limit", conf.count,
+                        "X-RateLimit-Remaining", 0,
+                        "X-RateLimit-Reset", reset)
+            end
+
             if conf.rejected_msg then
                 return conf.rejected_code, { error_msg = conf.rejected_msg }
             end
@@ -302,7 +309,8 @@ function _M.rate_limit(conf, ctx)
 
     if conf.show_limit_quota_header then
         core.response.set_header("X-RateLimit-Limit", conf.count,
-            "X-RateLimit-Remaining", remaining)
+                "X-RateLimit-Remaining", remaining,
+                "X-RateLimit-Reset", reset)
     end
 end
 

apisix/plugins/limit-count/limit-count-local.lua

@@ -0,0 +1,76 @@
+--
+-- 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_local_new = require("resty.limit.count").new
+local ngx = ngx
+local ngx_time = ngx.time
+local assert = assert
+local setmetatable = setmetatable
+
+local _M = {}
+
+local mt = {
+    __index = _M
+}
+
+local function set_endtime(self, key, time_window)
+    -- set an end time
+    local end_time = ngx_time() + time_window
+    -- save to dict by key
+    self.dict:set(key, end_time, time_window)
+
+    local reset = time_window
+    return reset
+end
+
+local function read_reset(self, key)
+    -- read from dict
+    local end_time = (self.dict:get(key) or 0)
+    local reset = end_time - ngx_time()
+    if reset < 0 then
+        reset = 0
+    end
+    return reset
+end
+
+function _M.new(plugin_name, limit, window, conf)
+    assert(limit > 0 and window > 0)
+
+    local self = {
+        limit_count = limit_local_new(plugin_name, limit, window, conf),
+        dict = ngx.shared["plugin-limit-count-reset-header"]
+    }
+
+    return setmetatable(self, mt)
+end
+
+function _M.incoming(self, key, commit, conf)
+    local delay, remaining = self.limit_count:incoming(key, commit)
+    local reset = 0
+    if not delay then
+        return delay, remaining, reset
+    end
+
+    if remaining == conf.count - 1 then
+        reset = set_endtime(self, key, conf.time_window)
+    else
+        reset = read_reset(self, key)
+    end
+
+    return delay, remaining, reset
+end
+
+return _M

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

@@ -30,11 +30,12 @@ local mt = {
 
 
 local script = core.string.compress_script([=[
-    if redis.call('ttl', KEYS[1]) < 0 then
+    local ttl = redis.call('ttl', KEYS[1])
+    if ttl < 0 then
         redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
-        return ARGV[1] - 1
+        return {ARGV[1] - 1, ARGV[2]}
     end
-    return redis.call('incrby', KEYS[1], -1)
+    return {redis.call('incrby', KEYS[1], -1), ttl}
 ]=])
 
 
@@ -91,16 +92,20 @@ function _M.incoming(self, key)
     local window = self.window
     key = self.plugin_name .. tostring(key)
 
-    local remaining, err = red:eval(script, 1, key, limit, window)
+    local ttl = 0
+    local res, err = red:eval(script, 1, key, limit, window)
 
     if err then
-        return nil, err
+        return nil, err, ttl
     end
 
+    local remaining = res[1]
+    ttl = res[2]
+
     if remaining < 0 then
-        return nil, "rejected"
+        return nil, "rejected", ttl
     end
-    return 0, remaining
+    return 0, remaining, ttl
 end
 
 

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

@@ -30,32 +30,17 @@ local mt = {
 
 
 local script = core.string.compress_script([=[
-    if redis.call('ttl', KEYS[1]) < 0 then
+    local ttl = redis.call('ttl', KEYS[1])
+    if ttl < 0 then
         redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
-        return ARGV[1] - 1
+        return {ARGV[1] - 1, ARGV[2]}
     end
-    return redis.call('incrby', KEYS[1], -1)
+    return {redis.call('incrby', KEYS[1], -1), ttl}
 ]=])
 
-
-function _M.new(plugin_name, limit, window, conf)
-    assert(limit > 0 and window > 0)
-
-    local self = {
-        limit = limit,
-        window = window,
-        conf = conf,
-        plugin_name = plugin_name,
-    }
-    return setmetatable(self, mt)
-end
-
-
-function _M.incoming(self, key)
-    local conf = self.conf
+local function redis_cli(conf)
     local red = redis_new()
     local timeout = conf.redis_timeout or 1000    -- 1sec
-    core.log.info("ttl key: ", key, " timeout: ", timeout)
 
     red:set_timeouts(timeout, timeout, timeout)
 
@@ -85,27 +70,52 @@ function _M.incoming(self, key)
         -- core.log.info(" err: ", err)
         return nil, err
     end
+    return red, nil
+end
+
+function _M.new(plugin_name, limit, window, conf)
+    assert(limit > 0 and window > 0)
+
+    local self = {
+        limit = limit,
+        window = window,
+        conf = conf,
+        plugin_name = plugin_name,
+    }
+    return setmetatable(self, mt)
+end
+
+function _M.incoming(self, key)
+    local conf = self.conf
+    local red, err = redis_cli(conf)
+    if err then
+        return red, err
+    end
 
     local limit = self.limit
     local window = self.window
-    local remaining
+    local res
     key = self.plugin_name .. tostring(key)
 
-    remaining, err = red:eval(script, 1, key, limit, window)
+    local ttl = 0
+    res, err = red:eval(script, 1, key, limit, window)
 
     if err then
-        return nil, err
+        return nil, err, ttl
     end
 
+    local remaining = res[1]
+    ttl = res[2]
+
     local ok, err = red:set_keepalive(10000, 100)
     if not ok then
-        return nil, err
+        return nil, err, ttl
     end
 
     if remaining < 0 then
-        return nil, "rejected"
+        return nil, "rejected", ttl
     end
-    return 0, remaining
+    return 0, remaining, ttl
 end
 
 

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

@@ -254,7 +254,7 @@ curl -i http://127.0.0.1:9180/apisix/admin/routes/1 \
 
 ## Example usage
 
-The above configuration limits to 2 requests in 60 seconds. The first two requests will work and the response headers will contain the headers `X-RateLimit-Limit` and `X-RateLimit-Remaining`:
+The above configuration limits to 2 requests in 60 seconds. The first two requests will work and the response headers will contain the headers `X-RateLimit-Limit` and `X-RateLimit-Remaining` and `X-RateLimit-Reset`, represents the total number of requests that are limited, the number of requests that can still be sent, and the number of seconds left for the counter to reset:
 
 ```shell
 curl -i http://127.0.0.1:9080/index.html
@@ -267,16 +267,20 @@ Content-Length: 13175
 Connection: keep-alive
 X-RateLimit-Limit: 2
 X-RateLimit-Remaining: 0
+X-RateLimit-Reset: 58
 Server: APISIX web server
 ```
 
-When you visit for a third time in the 60 seconds, you will receive a response with 503 code:
+When you visit for a third time in the 60 seconds, you will receive a response with 503 code. Currently, in the case of rejection, the limit count headers is also returned:
 
 ```shell
 HTTP/1.1 503 Service Temporarily Unavailable
 Content-Type: text/html
 Content-Length: 194
 Connection: keep-alive
+X-RateLimit-Limit: 2
+X-RateLimit-Remaining: 0
+X-RateLimit-Reset: 58
 Server: APISIX web server
 ```
 
@@ -287,6 +291,9 @@ HTTP/1.1 503 Service Temporarily Unavailable
 Content-Type: text/html
 Content-Length: 194
 Connection: keep-alive
+X-RateLimit-Limit: 2
+X-RateLimit-Remaining: 0
+X-RateLimit-Reset: 58
 Server: APISIX web server
 
 {"error_msg":"Requests are too frequent, please try again later."}

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

@@ -253,7 +253,7 @@ curl -i http://127.0.0.1:9180/apisix/admin/routes/1 \
 curl -i http://127.0.0.1:9080/index.html
 ```
 
-在执行测试命令的前两次都会正常访问。其中响应头中包含了 `X-RateLimit-Limit` 和 `X-RateLimit-Remaining` 字段，分别代表限制的总请求数和剩余还可以发送的请求数：
+在执行测试命令的前两次都会正常访问。其中响应头中包含了 `X-RateLimit-Limit` 和 `X-RateLimit-Remaining` 和 `X-RateLimit-Reset` 字段，分别代表限制的总请求数和剩余还可以发送的请求数以及计数器剩余重置的秒数：
 
 ```shell
 HTTP/1.1 200 OK
@@ -262,16 +262,20 @@ Content-Length: 13175
 Connection: keep-alive
 X-RateLimit-Limit: 2
 X-RateLimit-Remaining: 0
+X-RateLimit-Reset: 58
 Server: APISIX web server
 ```
 
-当第三次进行测试访问时，会收到包含 `503` HTTP 状态码的响应头，表示插件生效：
+当第三次进行测试访问时，会收到包含 `503` HTTP 状态码的响应头，目前在拒绝的情况下，也会返回相关的头，表示插件生效：
 
 ```shell
 HTTP/1.1 503 Service Temporarily Unavailable
 Content-Type: text/html
 Content-Length: 194
 Connection: keep-alive
+X-RateLimit-Limit: 2
+X-RateLimit-Remaining: 0
+X-RateLimit-Reset: 58
 Server: APISIX web server
 ```
 
@@ -282,6 +286,9 @@ HTTP/1.1 503 Service Temporarily Unavailable
 Content-Type: text/html
 Content-Length: 194
 Connection: keep-alive
+X-RateLimit-Limit: 2
+X-RateLimit-Remaining: 0
+X-RateLimit-Reset: 58
 Server: APISIX web server
 
 {"error_msg":"Requests are too frequent, please try again later."}

t/APISIX.pm

@@ -533,6 +533,7 @@ _EOC_
 
     lua_shared_dict plugin-limit-req 10m;
     lua_shared_dict plugin-limit-count 10m;
+    lua_shared_dict plugin-limit-count-reset-header 10m;
     lua_shared_dict plugin-limit-conn 10m;
     lua_shared_dict internal-status 10m;
     lua_shared_dict upstream-healthcheck 32m;