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

原理

令牌桶算法的原理是系统会以一个恒定的速度往桶里放入令牌，而如果请求需要被处理，则需要先从桶里获取一个令牌，当桶里没有令牌可取时，则我们可以选择等待也可以直接拒绝并返回。
令牌桶算法的优点是可以应对瞬时流量激增的情况，因为它可以通过调整发放令牌的速率来控制请求的处理速度。

实际令牌桶的实现中，并不用每隔一段时间更新当前桶的数量，而是记录了上次访问时和当前桶中令牌的数量，当再次访问时，通过上次访问时间计算出当前令牌的数量，决定是否可以发放。

1. 首先判断当前令牌桶是否存在，可能是过期或刚启动
2. 如果不存在就创建桶
3. 存在就通过时间计算差计算当前剩余值减一放回，更新时间。
4. 如果桶中剩余数量小于零，返回限流

使用

因为redis为了防止出现数据不一致的情况，不允许随机写入操作，因此选择将当前时间传入，而不是选择redis.call('time')

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

lua脚本：

-- 定义返回值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=KEYS[2]
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