APISIX源码解析-插件-固定窗口限速【limit-count】

limit-count 固定窗口限速插件

关键属性

源码实现

local function create_limit_obj(conf)
    core.log.info("create new limit-count plugin 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.access(conf, ctx)
    core.log.info("ver: ", ctx.conf_version)
    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)
        if conf.allow_degradation then
            return
        end
        return 500
    end

    local key = (ctx.var[conf.key] or "") .. 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
            if conf.rejected_msg then
                return conf.rejected_code, { error_msg = conf.rejected_msg }
            end
            return conf.rejected_code
        end

        core.log.error("failed to limit count: ", err)
        if conf.allow_degradation then
            return
        end
        return 500, {error_msg = "failed to limit count"}
    end

    if conf.show_limit_quota_header then
        core.response.set_header("X-RateLimit-Limit", conf.count,
            "X-RateLimit-Remaining", remaining)
    end
end

使用Openresty中的resty.limit.count模块

resty.limit.count 模块就是限制接口单位时间的请求数，

return limit_local_new("plugin-" .. plugin_name, conf.count,
                               conf.time_window)

该 new 方法携带的参数如下：
shdict_name： lua_shared_dict 声明的共享内存的名称。建议对不同的限制使用独立的共享内存。
count：指定的请求阈值。
time_window: 请求个数复位前的窗口时间，以秒为单位。

local delay, remaining = lim:incoming(key, true)

key: 是用户指定限制速率的 key。
commit：布尔值。如果设置为 true，则 lim将会在支持该 lim的共享内存中记录该事件；否则仅仅是 “dry run”。

返回值:

1、如果请求数没有超过在 new 方法中设置的 count 值，那么该 incoming 返回 0 作为
delay，remaining表示当前时间内余下允许请求的个数。
2、如果请求数超过了 count 限制，则返回 nil 和错误字符串 “rejected”。
3、如果发生错误（如访问共享内存失败），则该方法返回 nil 和错误描述字符串。

    if conf.show_limit_quota_header then
        core.response.set_header("X-RateLimit-Limit", conf.count,
            "X-RateLimit-Remaining", remaining)
    end

是否在响应头中显示 X-RateLimit-Limit 和 X-RateLimit-Remaining （限制的总请求数和剩余还可以发送的请求数），默认值是 true。

    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

同时支持通过redis或redis集群方式进行记录和限制，使用redis eval命令。
redis Eval 命令基本语法如下：
EVAL script numkeys key [key …] arg [arg …]
参数说明：
1、script： 参数是一段 Lua 5.1 脚本程序。脚本不必(也不应该)定义为一个 Lua 函数。
2、numkeys： 用于指定键名参数的个数。
3、key [key …]： 从 EVAL 的第三个参数开始算起，表示在脚本中所用到的那些 Redis 键(key)，这些键名参数可以在 Lua 中通过全局变量 KEYS 数组，用 1 为基址的形式访问( KEYS[1] ， KEYS[2] ，以此类推)。
4、arg [arg …]： 附加参数，在 Lua 中通过全局变量 ARGV 数组访问，访问的形式和 KEYS 变量类似( ARGV[1] 、 ARGV[2] ，诸如此类)。

function _M.incoming(self, key)
    ...
    local remaining, err = red:eval(script, 1, key, limit, window)
    ...
    
local script = [=[
    if redis.call('ttl', KEYS[1]) < 0 then
        redis.call('set', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2])
        return ARGV[1] - 1
    end
    return redis.call('incrby', KEYS[1], -1)
]=]

Redis命令操作
call :
脚本命令执行
ttl：
当 key 不存在时，返回 -2 。 当 key 存在但没有设置剩余生存时间时，返回 -1 。 否则，以秒为单位，返回 key 的剩余生存时间。
注意：在 Redis 2.8 以前，当 key 不存在，或者 key 没有设置剩余生存时间时，命令都返回 -1 。
set ex：
设置键值的同时，设置键的过期时间为 second 秒。 SET key value EX second 效果等同于 SETEX key second value 。
Incrby ：
将 key 中储存的数字加上指定的增量值。
如果 key 不存在，那么 key 的值会先被初始化为 0 ，然后再执行 INCRBY 命令。
如果值包含错误的类型，或字符串类型的值不能表示为数字，那么返回一个错误。
操作的值限制在 64 位(bit)有符号数字表示之内。

综上，Redis实现方式为：

1、当key不存在时，第一次访问，设置key : limit-1，key的过期时间为window，通过key过期机制，实现指定窗口时间内的计数重置
2、当key存在时，每次访问自动-1，代表剩余个数-1
3、当剩余个数 < 0 时，返回"rejected"，代表请求拒绝