OpenResty之 limit.count 模块

原文： lua-resty-limit-traffic/lib/resty/limit/count.md

1. 示例

http {
    lua_shared_dict my_limit_count_store 100m;
    
    init_by_lua_block {
        require "resty.core"
    }
    
    server {
        location / {
            access_by_lua_block {
                local limit_count = require "resty.limit.count"
                
                -- rate: 5000 requests per 3600s
                local lim, err = limit_count.new("my_limit_count_store", 5000, 3600)
                if not lim then
                    ngx.log(ngx.ERR, "failed to instantiate a resty.limit.coutn object: ", err)
                    return ngx.exit(500)
                end
                
                -- use the Authorization header as the limiting key
                local key = ngx.req.get_headers()["Authorization"] or "public"
                local delay, err = lim:incoming(key, true)
                
                if not delay then
                    if err == "rejected" then
                        ngx.header["X-RateLimit-Limit"] = "5000"
                        ngx.header["X-RateLimit-Remaining"] = 0
                        return ngx.exit(503)
                    end
                    ngx.log(ngx.ERR, "failed to limit count: ", err)
                    return ngx.exit(500)
                end
                
                -- the 2nd return value holds the current remaing number
                -- of requests for the specified key.
                local remaining = err
                
                ngx.header["X-RateLimit-Limit"] = "5000"
                ngx.header["X-RateLimit-Remaining"] = remaining
            }
        }
    }
}

注： 该模块依赖 lua-resty-core，因此需要：

init_by_lua_block {
    require "resty.core"
}

2. 方法

2.1 new

syntax: obj, err = class.new(shdict_name, count, time_window)

实例化 class 的对象，该 class 通过 require "resty.limit.count" 返回。

该 new 方法携带的参数如下：

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

new 实现如下

local ngx_shared = ngx.shared
local setmetatable = setmetatable
local assert = assert

local _M = {
    _VERSION = '0.05'
}

local mt = {
    __index = _M
}

-- the "limit" argument controls number of request allowed in a time window.
-- time "window" argument controls the time window in seconds.
function _M.new(dict_name, limit, window)
    local dict = ngx_shared[dict_name]
    if not dict then
        return nil, "shared dict not found"
    end
    
    assert(limit> 0 and window > 0)
    
    local self = {
        dict = dict,
        limit = limit,
        window = window,
    }
    
    return setmetatable(self, mt)
end

2.2 incoming

syntax: delay, err = obj:incoming(key, commit)

触发新请求传入事件并计算当前请求对指定 key 所需的 delay（如果有的话），或者是否立即拒绝该请求。

该方法接受如下参数：

• key: 是用户指定限制速率的 key。
  例如，可以使用 host 或 server zone 作为 key，以便限制每个 host 的速率。此外，也可以使用 Authorization 头部值作为 key，以便可以为个人用户限制速率。
  注意该模块没有为该 key 加前缀或后缀来标志该 key，因此用户需要确保该 key 在 lua_shared_dict 共享内存中是唯一的。
• commit：布尔值。如果设置为 true，则 obj 将会在支持该 obj 的共享内存中记录该事件；否则仅仅是 "dry run"。

该 incoming 方法的放回值依赖如下情况：

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

incoming 实现如下

function _M.incoming(self, key, commit)
    local dict = self.dict
    local limit = self.limit
    local window = self.window
    
    local remaining, ok, err
    
    if commit then
        remaining, err = dict:incr(key, -1, limit)
        if not remaining then
            return nil, err
        end
        
        if remaining == limit - 1 then
            ok, err = dict:expire(key, window)
            if not ok then
                if err == "not found" then
                    remaining, err = dict:incr(key, -1, limit)
                    if not remaining then
                        return nil, err
                    end
                    
                    ok, err = dict:expire(key, window)
                    if not ok then
                        return nil, err
                    end
                    
                else
                    return nil, err
                end
            end
        end
        
    else
        remaining = (dict:get(key) or limit) - 1
    end
    
    if remaining < 0 then
        return nil, "rejected"
    end
    
    return 0, remaining
end