526互联

Redis+lua 实现令牌桶限流算法

发布时间 2023-04-27 19:26:52作者: Bruce13

使用

lua := redis.NewScript(script)
args[0] = strconv.Itoa(fillInterval)
args[1] = strconv.FormatInt(time.Now().Unix()*1000, 10)
res, err := lua.Run(context.Background(), utils.Red, []string{"RateLimit"}, args[0], args[1]).Result()

-- 定义返回值res[1]是否触发限流（1限流 0通过）res[2]当前桶中的令牌数
local res={}
res[1]=0
--local curtime=redis.call('time')
local inteval_time=tonumber(ARGV[1]) -- 放入令牌桶的间隔时间
local current_time=tonumber(ARGV[2]) -- 当前时间

local amount=10 -- 一次取几个
local key_expire_time=1000*3600 -- 过期时间
local inflow__per_unit=100 -- 每次放多少
local capacity=1000
local st_key='last_update'
local bucket_amount = 0

-- 上次向桶中投放令牌的时间
local last_time=redis.call('get',st_key)
-- 当前令牌数
local current_value = redis.call('get',KEYS[1])

if(last_time == false or current_value == false) -- 令牌桶也不存在或过期，重新生成令牌桶
then
    bucket_amount = capacity - amount;
    -- 生成令牌桶
    redis.call('set',KEYS[1],bucket_amount,'PX',key_expire_time)
    -- 设置投放时间
    redis.call('set',st_key,current_time,'PX',key_expire_time)
    res[2]=bucket_amount
    return res
end

current_value = tonumber(current_value)
last_time=tonumber(last_time)
local past_time=current_time-last_time --当前时间-上次投放的时间
if(past_time<inteval_time)
then
    -- 不到放入令牌时间，直接从令牌桶中取走令牌
    bucket_amount=current_value-amount
else
    -- 需要放入令牌
    local cur_times = past_time/inteval_time -- 放几次

    cur_times=math.floor(cur_times)

    bucket_amount=current_value+cur_times*inflow__per_unit
    if (bucket_amount > capacity)
    then
        bucket_amount = capacity-amount
    end
    -- 有新投放，更新投放时间
    redis.call('set',st_key,current_time,'PX',key_expire_time)
end

res[2]=bucket_amount

-- 触发限流
if(bucket_amount<0)
then
    res[1]=1
    return res
end

-- 更新令牌桶KV
redis.call('set',KEYS[1],bucket_amount,'PX',key_expire_time)

return res