Redis Lua脚本实现令牌桶限流算法

脚本代码

-- 参数说明：
-- KEYS[1]: 令牌桶的Key（如 rate_limit:api1）
-- ARGV[1]: 令牌桶容量（最大令牌数）
-- ARGV[2]: 令牌生成速率（每秒生成数）
-- ARGV[3]: 当前时间戳（秒）
-- ARGV[4]: 请求令牌数（通常为1）

local key = KEYS[1]
local capacity = tonumber(ARGV[1])
local rate = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

-- 1. 获取桶内当前状态
local data = redis.call("HMGET", key, "tokens", "last_refill_time")
local tokens = tonumber(data[1]) or capacity  -- 初始化为满桶
local last_refill = tonumber(data[2]) or now

-- 2. 计算时间差并生成新令牌
local delta = math.max(now - last_refill, 0)
local new_tokens = math.floor(delta * rate)
tokens = math.min(tokens + new_tokens, capacity)  -- 不超过桶容量

-- 3. 更新最后一次补充时间（仅在生成新令牌时更新）
if new_tokens > 0 then
    redis.call("HSET", key, "last_refill_time", now)
end

-- 4. 判断是否有足够令牌
local result = 0
if tokens >= requested then
    tokens = tokens - requested
    result = 1  -- 允许请求
    redis.call("HSET", key, "tokens", tokens)
else
    -- 令牌不足，保留当前令牌数
    redis.call("HSET", key, "tokens", tokens)
end

-- 5. 设置Key过期时间（避免冷数据占用内存）
redis.call("EXPIRE", key, math.ceil(capacity / rate) * 2)

return {result, tokens}  -- 返回是否允许和剩余令牌数

---

脚本逐行解析

1. 参数接收

   • KEYS[1]：标识不同限流目标的唯一Key（如按API路径 + 用户ID）。
   • ARGV：传递容量、速率、当前时间戳和请求令牌数。

2. 初始化桶状态

   • 使用 HMGET 获取当前令牌数（tokens）和最后一次补充时间（last_refill_time）。
   • 若首次访问，默认令牌桶为满容量（tokens = capacity）。

3. 令牌补充逻辑

   • 计算自上次补充以来的时间差（delta），生成新令牌数 new_tokens = delta * rate。
   • 确保令牌总数不超过桶容量（math.min）。

4. 令牌消费判断

   • 若剩余令牌足够，扣除请求数并返回允许（result = 1）。
   • 若不足，保留当前令牌数（突发后可能需要时间恢复）。

5. 过期时间设置

   • 根据容量和速率计算合理TTL（如容量10，速率5/秒，TTL=4秒），避免内存泄漏。

     指令牌长时间没人拿然后一直占用内存