令牌桶算法实现方式

目录

一、令牌桶算法的核心原理

二、令牌桶的代码实现

1. 手动实现一个简单的令牌桶

2. 生产级方案：Guava RateLimiter

3. 分布式环境实现：Redis + Lua

三、关键设计点总结

令牌桶算法的实现可以分为基础原理和代码落地两个层面。我们先从它的核心工作机制说起，这样你就能理解为什么它能允许突发流量。

一、令牌桶算法的核心原理

令牌桶的工作过程就像是一个生产消费模型：

关键特性解析：

• 平滑速率：长期来看，请求速率被限制在令牌生成速率（比如10个/秒）附近。

• 突发流量：如果系统一段时间空闲，桶里会积攒最多20个令牌（容量值）。当突发流量到来时，只要桶里有20个令牌，就可以瞬间全部取走，允许这20个请求并发执行。这就是令牌桶能应对突发的原因。

二、令牌桶的代码实现

在实际代码中，我们不会真的开一个定时器不停放令牌，那样太消耗资源。标准的实现方式是延迟计算：在每次请求到来时，根据当前时间计算出这段时间应该生成多少令牌，然后一次性补充到桶里。

1. 手动实现一个简单的令牌桶

import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
* 手动实现一个简单的令牌桶 (线程安全)
*/
public class TokenBucket {
// 桶的容量
private final long capacity;
// 令牌生成速率 (每秒生成的令牌数)
private final long refillTokensPerSecond;
// 当前桶中的令牌数量 (使用AtomicLong保证线程安全)
private final AtomicLong currentTokens;
// 上次令牌补充的时间戳 (毫秒)
private volatile long lastRefillTimestamp;

/**
* @param capacity 桶容量
* @param refillTokensPerSecond 每秒生成的令牌数
*/
public TokenBucket(long capacity, long refillTokensPerSecond) {
this.capacity = capacity;
this.refillTokensPerSecond = refillTokensPerSecond;
// 初始时桶是满的，允许初始突发
this.currentTokens = new AtomicLong(capacity);
this.lastRefillTimestamp = System.currentTimeMillis();
}

/**
* 尝试获取一个令牌 (非阻塞)
* @return true: 获取成功, false: 获取失败
*/
public boolean tryAcquire() {
// 1. 补充令牌 (基于当前时间计算)
refill();

// 2. 尝试扣减令牌
while (true) {
long current = currentTokens.get();
if (current <= 0) {
return false; // 无可用令牌
}
if (currentTokens.compareAndSet(current, current – 1)) {
return true; // CAS扣减成功
}
// CAS失败说明有并发冲突，重试
}
}

/**
* 补充令牌 (核心逻辑：基于时间差的延迟计算)
*/
private void refill() {
long now = System.currentTimeMillis();
// 计算距离上次补充过去了多少毫秒
long timePassedMs = now – lastRefillTimestamp;
if (timePassedMs < 0) { // 处理时钟回拨
lastRefillTimestamp = now;
return;
}

// 根据时间计算应该生成多少新令牌
// 计算公式: (时间差 / 1000) * 每秒生成速率
long newTokens = (timePassedMs * refillTokensPerSecond) / 1000;

if (newTokens > 0) {
// 更新补充时间戳
lastRefillTimestamp = now;
// 增加令牌，但不能超过容量
while (true) {
long current = currentTokens.get();
long updated = Math.min(capacity, current + newTokens);
if (currentTokens.compareAndSet(current, updated)) {
break;
}
}
}
}

// 简单的测试
public static void main(String[] args) throws InterruptedException {
// 创建一个桶：容量5，每秒产生2个令牌
TokenBucket bucket = new TokenBucket(5, 2);

// 突发消费：瞬间把5个令牌全部取走
for (int i = 0; i < 5; i++) {
System.out.println("请求1-" + i + ": " + bucket.tryAcquire());
}
// 第6个请求会失败 (桶空了)
System.out.println("请求1-5: " + bucket.tryAcquire());

// 等待1秒 (会产生2个令牌)
Thread.sleep(1000);

// 可以成功获取2个
System.out.println("请求2-0: " + bucket.tryAcquire());
System.out.println("请求2-1: " + bucket.tryAcquire());
// 这一个会失败
System.out.println("请求2-2: " + bucket.tryAcquire());
}
}

实现要点说明：

• CAS（Compare and Swap，比较并交换）操作：通过AtomicLong和compareAndSet保证并发环境下的线程安全，避免加锁带来的性能开销。

• 时间差计算：通过计算当前时间与上次补充时间的差值，推算出需要补充的令牌数，避免了主动创建定时器线程。

• 容量限制：补充令牌时使用Math.min(capacity, current + newTokens)，确保令牌数不会超过桶的容量。

2. 生产级方案：Guava RateLimiter

在实际项目中，通常不需要重复造轮子，Google Guava 提供的 RateLimiter 是业界使用最广泛的令牌桶实现。

// 引入依赖
// <dependency>
// <groupId>com.google.guava</groupId>
// <artifactId>guava</artifactId>
// <version>32.1.3-jre</version>
// </dependency>

import com.google.common.util.concurrent.RateLimiter;

public class GuavaRateLimiterExample {
public static void main(String[] args) {
// 1. 创建限流器：每秒产生5个令牌 (平滑模式)
RateLimiter limiter = RateLimiter.create(5);

// 2. 使用tryAcquire非阻塞获取
if (limiter.tryAcquire()) {
System.out.println("获取令牌成功，执行业务逻辑");
} else {
System.out.println("限流中，返回错误或降级");
}

// 3. 使用acquire阻塞获取 (最多等待指定时间)
// limiter.acquire(); // 如果没有令牌，会一直阻塞直到获取到

// 4. 支持预热模式
// RateLimiter.create(5, 10, TimeUnit.SECONDS); // 每秒5个，但需要10秒预热到该速率
}
}

Guava RateLimiter 的两个核心子类：

• SmoothBursty (默认)：经典的令牌桶，允许突发流量。

• SmoothWarmingUp：带有预热的令牌桶。当系统长时间空闲后，突发流量不会立刻达到最大速率，而是平滑地逐步提升，避免冷启动时把数据库等下游资源打垮。

3. 分布式环境实现：Redis + Lua

在分布式系统中，需要将令牌状态存储在共享的Redis中。这里提供一个Lua脚本示例，它实现了原子性的令牌补充和扣减。

— Lua脚本: token_bucket.lua
— KEYS[1]: 令牌桶的key
— ARGV[1]: 桶容量
— ARGV[2]: 当前时间戳 (秒)
— ARGV[3]: 令牌生成速率 (每秒)
— ARGV[4]: 请求的令牌数 (通常为1)

— 1. 获取当前桶中令牌数 (如果没有则初始化为容量)
local current = redis.call('get', KEYS[1])
local last_refill_time = redis.call('get', KEYS[1]..':time')

if current == false then
current = ARGV[1] — 初始时令牌满
else
— 2. 计算时间差，补充令牌 (关键逻辑)
local time_diff = tonumber(ARGV[2]) – tonumber(last_refill_time or ARGV[2])
local new_tokens = time_diff * tonumber(ARGV[3])
if new_tokens > 0 then
current = math.min(tonumber(ARGV[1]), tonumber(current) + new_tokens)
else
current = tonumber(current)
end
end

— 3. 判断是否足够
if current >= tonumber(ARGV[4]) then
current = current – tonumber(ARGV[4])
redis.call('set', KEYS[1], current)
redis.call('set', KEYS[1]..':time', ARGV[2])
return 1 — 表示成功
else
— 注意：即使令牌不足，也需要更新时间，防止把时间差累积到下次
— 这里简化为只返回失败，实际也可只更新时间不扣减
redis.call('set', KEYS[1]..':time', ARGV[2])
return 0 — 表示被限流
end

调用方式： Java端通过redisTemplate.execute(script, keys, args)调用此脚本。由于整个逻辑在Redis中原子执行，完美支持分布式环境下的并发限流。

三、关键设计点总结