分布式限流——Redis + Lua脚本实现令牌桶算法

主要思路概括如下：

定义数据结构：
- 使用Redis存储令牌桶的状态，包括当前令牌数（KEYS $1$ ）和上一次令牌填充的时间戳（KEYS $1$ :last）。
计算新增令牌：
- 获取当前系统时间与上次令牌填充时间的时间差，并基于令牌生成速率计算在这段时间内应新增的令牌数。
- 确保新增令牌数不超过桶的总容量。
更新令牌数：
- 将令牌桶内的令牌数增加至新的值，确保不超过桶的最大容量。
判断是否满足请求：
- 如果更新后的令牌数足以满足本次请求消耗的令牌数，则扣除相应数量的令牌并返回1，表示请求通过。
- 否则，返回0，表示请求被限流。
原子性操作：
- 利用Redis Lua脚本提供的原子性执行能力，确保在多客户端并发访问时，令牌桶状态的读取、更新等操作是线程安全的。
实时性：
- 脚本设计考虑到随着时间推移动态补充令牌，但实际应用中可能需要更精确的定时任务或者使用Redis的过期时间特性来定期补充令牌。

整个令牌桶算法在Redis和Lua脚本中的实现旨在提供一种简单且高效的手段来进行流量控制，既可以应对突发流量，又能保持总体速率稳定可控。

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命令所提供的功能。