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