主要思路概括如下:
-
定义数据结构:
- 使用Redis存储令牌桶的状态,包括当前令牌数(KEYS[1])和上一次令牌填充的时间戳(KEYS[1]:last)。
-
计算新增令牌:
- 获取当前系统时间与上次令牌填充时间的时间差,并基于令牌生成速率计算在这段时间内应新增的令牌数。
- 确保新增令牌数不超过桶的总容量。
-
更新令牌数:
- 将令牌桶内的令牌数增加至新的值,确保不超过桶的最大容量。
-
判断是否满足请求:
- 如果更新后的令牌数足以满足本次请求消耗的令牌数,则扣除相应数量的令牌并返回1,表示请求通过。
- 否则,返回0,表示请求被限流。
-
原子性操作:
- 利用Redis Lua脚本提供的原子性执行能力,确保在多客户端并发访问时,令牌桶状态的读取、更新等操作是线程安全的。
-
实时性:
- 脚本设计考虑到随着时间推移动态补充令牌,但实际应用中可能需要更精确的定时任务或者使用Redis的过期时间特性来定期补充令牌。
整个令牌桶算法在Redis和Lua脚本中的实现旨在提供一种简单且高效的手段来进行流量控制,既可以应对突发流量,又能保持总体速率稳定可控。
-- KEYS[1] 是令牌桶的键名
-- ARGV[1] 是令牌生成速率(每秒产生的令牌数量)
-- ARGV[2] 是桶容量(最大令牌存储量)
-- ARGV[3] 是请求需要消耗的令牌数量
-- ARGV[4] 是当前时间戳(毫秒级,由客户端提供)
-- 获取当前桶内令牌数
local tokenCount = tonumber(redis.call('GET', KEYS[1]))
-- 如果桶内令牌数为空或不存在,则初始化为桶容量
if tokenCount == nil then
tokenCount = tonumber(ARGV[2])
end
-- 计算从上次填充到现在应该产生的令牌数
local currentTime = tonumber(ARGV[4])
local lastRefillTimestamp = tonumber(redis.call('GET', KEYS[1] .. ':last'))
if lastRefillTimestamp == nil then
lastRefillTimestamp = currentTime
end
local newTokens = math.min((currentTime - lastRefillTimestamp) * tonumber(ARGV[1]), tonumber(ARGV[2]) - tokenCount)
-- 更新令牌数(注意这里假设时间粒度较小,不考虑精确到毫秒的令牌生成)
tokenCount = math.min(tokenCount + newTokens, tonumber(ARGV[2]))
-- 设置最后填充令牌的时间为当前时间
redis.call('SET', KEYS[1] .. ':last', currentTime)
-- 如果新加入的令牌仍然不够消耗,则请求被限流
if tokenCount < tonumber(ARGV[3]) then
return 0 -- 表示限流
else
-- 消耗相应数量的令牌
tokenCount = tokenCount - tonumber(ARGV[3])
redis.call('SET', KEYS[1], tokenCount)
return 1 -- 表示通过
end
在上面的Lua脚本示例中,上次填充令牌的时间戳是通过另一个Redis键来保存的。在脚本中,这个键是通过在令牌桶的主键后面加上:last_refill
来构成的,如 KEYS[1] .. ":last_refill"
。
这意味着对于一个令牌桶的键 bucket_key
,它的最后一次填充令牌的时间戳会被保存在 bucket_key:last_refill
这个额外的键中。
在脚本执行过程中,当需要获取上次填充令牌的时间戳时,通过调用 redis.call('GET', KEYS[1] .. ":last_refill")
来获取。如果没有获取到值(即第一次运行或者之前没有设置过),则会将其默认设置为当前时间戳。
同样,在完成令牌填充的逻辑后,也会更新这个时间戳键的值,通过 redis.call('SET', KEYS[1] .. ":last_refill", currentTime)
实现,其中 currentTime
是客户端提供的当前时间戳。
这个Lua脚本简化了令牌桶算法的实现,实际应用中可能还需要考虑更多细节,例如如何精确地按照时间间隔补充令牌、如何处理多个请求同时竞争令牌等情况。客户端在调用这个Lua脚本时,应当传递正确的参数,包括令牌桶的键名以及令牌相关的配置信息。此外,在高并发场景下,为了避免多个客户端同时修改令牌数,Lua脚本的执行必须是原子性的,这正是Redis.eval/evalsha命令所提供的功能。