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
posted @ 2018-08-06 22:56  季末的天堂  阅读(1586)  评论(0编辑  收藏