xuj789 / distributed-current-limiter Goto Github PK

--- 判断source_str 中是否contains sub_str
--- @param source_str
--- @param sub_str
local function contains(source_str , sub_str)
    local start_pos,end_pos = string.find(source_str,sub_str);
    if start_pos == nil then
        return false
    end
    local source_str_len = string.len(source_str)
    local s = string.sub(source_str,end_pos,end_pos)
    if source_str_len == end_pos then
        return true;
    elseif string.sub(source_str,end_pos + 1,end_pos + 1) == "," then
        return true;
    end
    return false;
end


--- 获取令牌
--- 返回码
--- 0 没有令牌桶配置
--- -1 表示取令牌失败，也就是桶里没有令牌
--- 1 表示取令牌成功
--- @param key 令牌的唯一标识
--- @param permits  请求令牌数量
--- @param curr_mill_second 当前毫秒数
--- @param context 使用令牌的应用标识
local function acquire(key , permits , curr_mill_second , context)
    local rate_limit_info = redis.pcall("HMGET" , key , "last_mill_second","curr_permits","max_permits","rate","apps")
    local last_mill_second = rate_limit_info[1]
    local curr_permits = tonumber(rate_limit_info[2])
    local max_permits = tonumber(rate_limit_info[3])
    local rate = rate_limit_info[4]
    local apps = rate_limit_info[5]

    -- 标识没有配置令牌桶
    if type(apps) == 'boolean' or apps == nil or not contains(apps , context) then
        return 0
    end

    local local_curr_permits = curr_permits


    --- 令牌桶刚刚创建，上一次获取令牌的毫秒数为空
    --- 根据和上一次向桶里添加令牌的时间和当前时间差，触发式往桶里添加令牌，并且更新上一次向桶里添加令牌的时间
    --- 如果向桶里添加的令牌数不足一个，则不更新上一次向桶里添加令牌的时间
    if (type(last_mill_second) ~= 'boolean' and last_mill_second ~= nil) then
        local reverse_permits = math.floor(((curr_mill_second - last_mill_second)/1000) * rate)
        local expect_curr_permits = reverse_permits + curr_permits;
        local_curr_permits = math.min(expect_curr_permits,max_permits)
        --- 大于0表示不是第一次获取令牌，也没有向桶里添加令牌
        if (reverse_permits > 0) then
            redis.pcall("HSET",key, "last_mill_second" , curr_mill_second)
        end
    else
        redis.pcall("HSET" , key, "last_mill_second" , curr_mill_second)
    end
    local result = -1
    if (local_curr_permits - permits >=0) then
        result = 1
        redis.pcall("HSET" , key , "curr_permits" , local_curr_permits - permits)
    else
        redis.pcall("HSET" , key , "curr_permits" , local_curr_permits)
    end
    return result
end


--- 初始化令牌桶配置
--- @param key 令牌的唯一标识
--- @param max_permits 桶大小
--- @param rate  向桶里添加令牌的速率
--- @param apps  可以使用令牌桶的应用列表，应用之前用逗号分隔
local function init(key, max_permits, rate , apps)
    local rate_limit_info = redis.pcall("HMGET",key,"last_mill_second","curr_permits","max_permits","rate","apps")
    local org_max_permits = tonumber(rate_limit_info[3])
    local org_rate = rate_limit_info[4]
    local org_apps = rate_limit_info[5]
    if (org_max_permits == nil) or (apps ~= org_apps or rate ~= org_rate or max_permits ~= org_max_permits) then
        redis.pcall("HMSET", key , "max_permits" , max_permits , "rate" , rate , "curr_permits" , max_permits , "apps" , apps)
    end
    return 1
end

--- 删除令牌桶
local function delete(key)
    redis.pcall("DEL",key)
    return 1
end


local key = KEYS[1]
local method = ARGV[1]
if method == 'acquire' then
    return acquire(key,KEYS[2],KEYS[3],KEYS[4])
elseif method == 'init' then
    return init(key,KEYS[2],KEYS[3],KEYS[4])
elseif method == 'delete' then
    return delete(key)
end