令牌桶算法Java实现代码解析
令牌桶算法是一种常用的流量控制/限流算法，下面我将详细解释前面提供的几种Java实现方案。
方案1：Guava RateLimiter实现
Guava是Google提供的Java核心库，其中RateLimiter类实现了令牌桶算法。
核心代码：
import com.google.common.util.concurrent.RateLimiter;
public class GuavaTokenBucket {
private final RateLimiter rateLimiter;
public GuavaTokenBucket(double permitsPerSecond) {
this.rateLimiter = RateLimiter.create(permitsPerSecond);
}
public boolean tryAcquire() {
return rateLimiter.tryAcquire();
}
public boolean tryAcquire(int permits) {
return rateLimiter.tryAcquire(permits);
}
}
解析：
- 初始化：
RateLimiter.create(permitsPerSecond)创建一个每秒产生指定数量令牌的限流器 - 获取令牌：
tryAcquire()尝试获取1个令牌，立即返回成功或失败tryAcquire(int permits)尝试获取多个令牌
- 特点：
- 线程安全
- 支持预热期配置
- 支持突发流量处理
- 内部实现考虑了高精度时间计算
方案2：自定义基本实现
核心代码解析：
public class TokenBucket {
private final long capacity; // 桶的最大容量
private final long refillInterval; // 填充间隔(毫秒)
private final long refillTokens; // 每次填充的令牌数
private AtomicLong availableTokens; // 当前可用令牌数
private long lastRefillTimestamp; // 上次填充时间戳
public TokenBucket(long capacity, long refillTokens, long refillInterval, TimeUnit timeUnit) {
// 初始化参数转换
this.capacity = capacity;
this.refillTokens = refillTokens;
this.refillInterval = timeUnit.toMillis(refillInterval);
this.availableTokens = new AtomicLong(0);
this.lastRefillTimestamp = System.currentTimeMillis();
}
public synchronized boolean tryAcquire(int tokens) {
refill(); // 先补充令牌
if (availableTokens.get() >= tokens) {
availableTokens.addAndGet(-tokens);
return true;
}
return false;
}
private void refill() {
long now = System.currentTimeMillis();
long elapsedTime = now - lastRefillTimestamp;
// 计算需要填充的令牌批次数
long refillCount = elapsedTime / refillInterval;
if (refillCount > 0) {
// 计算新令牌数(不超过容量)
long newTokens = Math.min(capacity,
availableTokens.get() + refillCount * refillTokens);
availableTokens.set(newTokens);
lastRefillTimestamp = now;
}
}
}
关键点解析：
-
参数说明：
capacity：桶的最大容量，防止令牌无限累积refillInterval和refillTokens：每间隔X时间单位补充Y个令牌availableTokens：当前可用的令牌数(使用AtomicLong保证原子性)lastRefillTimestamp：上次补充令牌的时间戳
-
令牌补充逻辑(refill方法)：
- 计算自上次补充后经过的时间
- 计算应该补充的令牌批次(整数次)
- 补充令牌但不超过最大容量
-
获取令牌逻辑(tryAcquire方法)：
- 使用
synchronized保证线程安全 - 先补充令牌
- 检查是否有足够令牌
- 有则扣除并返回true，否则返回false
- 使用
方案3：高效自定义实现
改进点解析：
public class EfficientTokenBucket {
private final long capacity;
private final double refillTokensPerMillis; // 每毫秒补充的令牌数
private AtomicLong lastRefillTime;
private AtomicLong availableTokens;
public EfficientTokenBucket(long capacity, long refillTokens, long refillInterval, TimeUnit timeUnit) {
this.capacity = capacity;
// 计算每毫秒的令牌补充速率
this.refillTokensPerMillis = (double) refillTokens / timeUnit.toMillis(refillInterval);
this.lastRefillTime = new AtomicLong(System.currentTimeMillis());
this.availableTokens = new AtomicLong(0);
}
public boolean tryAcquire(long tokens) {
refill();
long currentAvailable;
long newAvailable;
do {
currentAvailable = availableTokens.get();
if (currentAvailable < tokens) {
return false;
}
newAvailable = currentAvailable - tokens;
} while (!availableTokens.compareAndSet(currentAvailable, newAvailable));
return true;
}
private void refill() {
long now = System.currentTimeMillis();
long lastTime = lastRefillTime.get();
long elapsedTime = now - lastTime;
if (elapsedTime > 0) {
double newTokens = elapsedTime * refillTokensPerMillis;
long newAvailable = Math.min(capacity,
availableTokens.get() + (long) newTokens);
if (lastRefillTime.compareAndSet(lastTime, now)) {
availableTokens.set(newAvailable);
}
}
}
}
任务调度线程池实现
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
public class TokenBucket {
private final int capacity; // 令牌桶的容量
private final AtomicInteger tokens; // 当前令牌数量
private final ScheduledExecutorService scheduler;
public TokenBucket(int capacity, int refillRate) {
this.capacity = capacity;
this.tokens = new AtomicInteger(0);
this.scheduler = Executors.newScheduledThreadPool(1);
// 定期添加令牌
scheduler.scheduleAtFixedRate(() -> {
int currentTokens = tokens.get();
int newTokens = Math.min(capacity, currentTokens + refillRate);
tokens.set(newTokens);
}, 0, 1, TimeUnit.SECONDS); // 每秒添加 refillRate 个令牌
}
public boolean tryConsume() {
while (true) {
int currentTokens = tokens.get();
if (currentTokens <= 0) {
return false; // 令牌不足，请求被拒绝
}
if (tokens.compareAndSet(currentTokens, currentTokens - 1)) {
return true; // 成功获取令牌，请求通过
}
}
}
public void shutdown() {
scheduler.shutdown();
}
public static void main(String[] args) {
TokenBucket tokenBucket = new TokenBucket(10, 2); // 容量为10，每秒补充2个令牌
// 模拟请求
for (int i = 0; i < 20; i++) {
if (tokenBucket.tryConsume()) {
System.out.println("Request " + i + " processed.");
} else {
System.out.println("Request " + i + " rejected.");
}
try {
Thread.sleep(500); // 模拟请求间隔
} catch (InterruptedException e) {
e.printStackTrace();
}
}
tokenBucket.shutdown();
}
}
代码说明 TokenBucket 类： capacity：令牌桶的最大容量。 tokens：当前令牌数量，使用 AtomicInteger 保证线程安全。 scheduler：用于定期向令牌桶中添加令牌的调度器。 构造函数： 初始化令牌桶的容量和当前令牌数量。 使用 ScheduledExecutorService 每秒向令牌桶中添加 refillRate 个令牌。 tryConsume 方法： 尝试从令牌桶中获取一个令牌。 如果当前令牌数量大于0，则通过 CAS 操作减少令牌数量并返回 true。 如果令牌数量不足，则返回 false。 shutdown 方法： 关闭调度器，停止向令牌桶中添加令牌。 main 方法： 创建一个令牌桶实例，容量为10，每秒补充2个令牌。 模拟20个请求，每个请求间隔500毫秒。 根据 tryConsume 的返回值判断请求是否被处理或拒绝。 通过这种方式，可以有效地控制请求的速率，防止系统过载。
改进点说明：
-
更精确的令牌补充：
- 使用
refillTokensPerMillis(每毫秒令牌数)代替批次补充 - 计算更精确，避免批次补充导致的不精确问题
- 使用
-
更高效的并发控制：
- 使用
compareAndSet(CAS)操作代替synchronized - 减少线程阻塞，提高并发性能
- 使用
-
原子变量优化：
- 使用
AtomicLong的原子操作保证线程安全 lastRefillTime也使用原子变量，避免单独使用时间戳可能的问题
- 使用
三种方案对比
| 特性 | Guava实现 | 基础自定义实现 | 高效自定义实现 |
|---|---|---|---|
| 线程安全 | 是 | 是(synchronized) | 是(CAS) |
| 精确度 | 高 | 中等(批次补充) | 高(连续补充) |
| 性能 | 高 | 中等 | 高 |
| 灵活性 | 中(配置选项有限) | 高 | 高 |
| 实现复杂度 | 低(直接使用) | 中 | 高 |
| 适用场景 | 大多数通用场景 | 简单需求 | 高性能需求 |
使用建议
- 优先考虑Guava实现：除非有特殊需求，否则Guava的实现已经足够好
- 需要精确控制时：选择高效自定义实现
- 简单场景：基础自定义实现足够
- 分布式环境：这些实现都是单机的，分布式系统需要Redis等分布式方案
每种实现都有其适用场景，选择时应根据具体需求权衡性能、精确度和实现复杂度。
