redis lua脚本令牌桶限流实现

redis.replicate_commands()

local key = KEYS[1] -- 令牌桶标识
local capacity = tonumber(ARGV[1]) -- 最大容量
local quota = tonumber(ARGV[2]) -- 时间窗口内的限额
local period = tonumber(ARGV[3]) -- 时间窗口大小（秒）
local quantity = tonumber(ARGV[4]) or 1 -- 需要的令牌数量，默认为1
local timestamp = tonumber(redis.call('time')[1]) -- 当前时间戳

-- 第一次请求时创建令牌桶
if (redis.call('exists', key) == 0) then
    redis.call('hmset', key, 'remain', capacity, 'timestamp', timestamp)
else
    -- 计算从上次生成到现在这段时间应该生成的令牌数
    local remain = tonumber(redis.call('hget', key, 'remain'))
    local last_reset = tonumber(redis.call('hget', key, 'timestamp'))
    local delta_quota = math.floor(((timestamp - last_reset) / period) * quota)
    if (delta_quota > 0) then
        remain = remain + delta_quota
        if (remain > capacity) then
            remain = capacity
        end
        redis.call('hmset', key, 'remain', remain, 'timestamp', timestamp)
    end
end

-- 支持动态调整容量和令牌生成速率
redis.call('hmset', key, 'capacity', capacity, 'quota', quota, 'period', period);

local result = {} -- 返回的结果集
local remain = tonumber(redis.call('hget', key, 'remain'))
if (remain < quantity) then
    result = {1, capacity, remain}
else
    result = {0, capacity, remain - quantity}
    redis.call('hincrby', key, 'remain', -quantity)
end

return result