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