Redis：Token Bucket 令牌桶算法

介绍常用的 Redis Token Bucket 令牌桶算法

Intro

在众多的限流与流量整形方案中，令牌桶算法（Token Bucket Algorithm）几乎是最常见的。其灵感源自于分组交换网络中的流量控制机制。让Gemini生成了一个科普小漫画：

从抽象层面来看，该算法通过维护一个固定容量的容器（即“桶”）来控制请求的通过率。其运作机制包含两个并发的过程：令牌的生产与令牌的消耗。

系统以一个恒定的速率（设为 $r$ ，单位为 tokens/sec）向桶中投放令牌。这个速率代表了系统允许的平均处理能力。与此同时，桶具有一个固定的容量上限（设为 $b$ ）。这一上限至关重要，它决定了系统能够容忍的最大突发量。当桶满时，新产生的令牌会被丢弃。

当一个请求（或数据包）到达时，它必须尝试从桶中获取一定数量的令牌（通常为 1 个，但在某些基于字节大小限流的场景下，可能对应数据包的大小）。

如果桶内令牌充足，请求获取令牌后将被放行，同时桶内令牌数扣减；
如果桶内令牌不足，请求则会被拒绝或进入队列等待，具体取决于限流策略是流量监管还是流量整形。

在任意时间段 $T$ 内，令牌桶算法允许通过的最大请求量 $N$ 遵循公式：

N = b + r \times T

令牌桶算法允许在短时间内处理的数据量超过平均速率 $r$ 。只要桶内有累积的令牌，系统就能以极快的速度（受限于物理带宽或处理性能）瞬间处理这批请求，直到桶被排空，随后处理速率将回归并被锁定在恒定速率 $r$ 上。

Redis实现Token Bucket

在分布式环境下实现令牌桶，最大的敌人是并发竞争（Race Condition）。如果两个服务实例同时读取 Redis 中的令牌数，发现都还剩 1 个，于是都决定“放行”并将令牌数减为 0 写回。结果是，原本只能通过 1 个请求，实际上却通过了 2 个。

为了解决这个问题，我们必须利用 Redis 的 Lua 脚本功能。Redis 保证 Lua 脚本在执行过程中是以原子方式运行的（其实就是Redis 是单线程处理命令的），这意味着脚本执行期间，不会有其他 Redis 命令插入执行。这相当于在 Redis 服务端加了一把锁。

Lua脚本

Lua是一种轻量级、高效的脚本语言。Lua 的解释器非常小（只有几百 KB），编译后的执行速度极快。在用户想要对redis执行一系列逻辑时，可以把这些逻辑写成一个Lua脚本发给 Redis。Redis 会将整个脚本作为一个原子操作执行。执行期间，其他任何命令都不会被插入。

执行流程

STEP0：定义参数

-- 参数说明：
-- KEYS[1]: 限流的唯一 Key (例如 rate_limit:api:user_123)
-- ARGV[1]: 桶的容量 (capacity)
-- ARGV[2]: 令牌生成速率 (rate, token/sec)
-- ARGV[3]: 当前请求需要的令牌数 (requested_tokens, 默认 1)

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

STEP1：获取时间，调用 Redis 内置的 TIME 命令获取时间戳：

1 2	local time_info = redis.call('TIME') local now = tonumber(time_info[1])

在分布式系统中，应用服务器可能有几十台，它们的系统时间很难做到绝对同步，因此，从redis获取时间是最稳妥的做法。

STEP2：获取当前桶的状态：使用 HMGET 一次性取出当前桶里剩多少令牌 (tokens)，以及上一次在这个桶里放入令牌是什么时候 (last_refill_time)。

1
2
3

local info = redis.call('HMGET', key, 'tokens', 'last_refill_time')
local current_tokens = tonumber(info[1])
local last_refill_time = tonumber(info[2])

STEP3：初始化：判断 Redis 里有没有这个 Key。如果 current_tokens 为空，说明这是系统第一次遇到这个请求（或者 Key 过期了）。此时，我们直接把桶填满，并将上次填充时间设为现在。

if current_tokens == nil then
    current_tokens = capacity
    last_refill_time = now
end

STEP4：惰性填充：计算时间差 (delta)、应补令牌 (filled)、容量限制 (math.min)、更新时间。

if now > last_refill_time then
    local delta = now - last_refill_time
    local filled = delta * rate
    current_tokens = math.min(capacity, current_tokens + filled)
    last_refill_time = now
end

STEP5：扣减令牌与决策：判断令牌是否足够支持流量通过。

local allowed = false
if current_tokens >= requested then
    current_tokens = current_tokens - requested
    allowed = true
end

STEP6：状态回写与生命周期管理：将计算后的最新令牌数和最新时间写回 Redis，给这个 Key 设置一个过期时间。

1
2
3

redis.call('HMSET', key, 'tokens', current_tokens, 'last_refill_time', last_refill_time)
redis.call('EXPIRE', key, math.ceil(capacity/rate) + 60)
return allowed

完整代码

-- rate_limit.lua
-- 参数说明：
-- KEYS[1]: 限流的唯一 Key (例如 rate_limit:api:user_123)
-- ARGV[1]: 桶的容量 (capacity)
-- ARGV[2]: 令牌生成速率 (rate, token/sec)
-- ARGV[3]: 当前请求需要的令牌数 (requested_tokens, 默认 1)

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

-- 获取 Redis 服务器当前时间（秒级和微秒级），保证分布式时钟一致性
-- 相比由客户端传入时间，使用 Redis 自身时间能避免服务器时钟不同步问题
local time_info = redis.call('TIME')
local now = tonumber(time_info[1])

-- 获取当前桶的剩余令牌数和上次填充时间
local info = redis.call('HMGET', key, 'tokens', 'last_refill_time')
local current_tokens = tonumber(info[1])
local last_refill_time = tonumber(info[2])

-- 初始化：如果 key 不存在，说明是第一次访问，默认桶是满的
if current_tokens == nil then
    current_tokens = capacity
    last_refill_time = now
end

-- 计算惰性填充
if now > last_refill_time then
    local delta = now - last_refill_time
    -- 计算这段时间生成的令牌：时间差 * 速率
    local filled = delta * rate
    -- 令牌数不能超过桶容量
    current_tokens = math.min(capacity, current_tokens + filled)
    -- 更新填充时间
    last_refill_time = now
end

-- 判断令牌是否足够
local allowed = false
if current_tokens >= requested then
    current_tokens = current_tokens - requested
    allowed = true
end

-- 更新 Redis 状态
redis.call('HMSET', key, 'tokens', current_tokens, 'last_refill_time', last_refill_time)
-- 设置 Key 过期时间（例如 60 秒），防止冷数据长期占用内存
-- 过期时间应至少大于 (capacity / rate)，确保桶满之前不会丢失数据
redis.call('EXPIRE', key, math.ceil(capacity/rate) + 60)

return allowed

SpringBoot 使用 Token Bucket

STEP0：将 Lua 脚本文件放在 src/main/resources/lua/rate_limit.lua

STEP1：在原本的配置类中通过 DefaultRedisScript 加载它，这样脚本会被预加载到 Redis 中。

@Configuration
public class RedisScriptConfig {
    @Bean
    public DefaultRedisScript<Long> limitScript() {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        // 设置脚本路径
        redisScript.setLocation(new ClassPathResource("lua/rate_limit.lua"));
        // 设置返回值类型
        redisScript.setResultType(Long.class);
        return redisScript;
    }
}

STEP2：编写调用工具类，写一个Service，封装调用逻辑：

@Service
public class RateLimiterService {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    @Autowired
    private DefaultRedisScript<Long> limitScript;

    public boolean shouldAllow(String key, int capacity, int rate) {
        // 参数说明：脚本, Key 列表, 参数列表
        Long result = stringRedisTemplate.execute(
                limitScript,
                Collections.singletonList(key),
                String.valueOf(capacity),
                String.valueOf(rate),
                "1" // 默认请求 1 个令牌
        );

        return result != null && result == 1L;
    }
}

STEP4：编写一个测试，这里我们引入泊松过程模拟，来模仿真实的流量访问，将最大容量设置为100，加入了集中大批量的访问场景，最终模拟结果如下图。

其中，绿色部分节点为成功访问，而红色的为拒绝访问。可以看到，在刚开始的稳定访问阶段，Token Bucket 的消耗与补充基本处于平衡状态，而在随之而来的集中访问阶段，token 的补充速度远远小于消耗速度，因此，导致了访问失败的情况，从而减少对服务器的流量冲击。