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