本文目录导读：
我来提供一个完整的Java令牌桶限流实现,包含简单版本和进阶版本。
基础令牌桶实现
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
/**
* 基础令牌桶限流器
*/
public class TokenBucket {
// 令牌桶容量
private final long capacity;
// 令牌产生速率（每秒）
private final double refillRate;
// 当前令牌数
private double tokens;
// 上次补充令牌的时间
private long lastRefillTime;
public TokenBucket(long capacity, double refillRate) {
this.capacity = capacity;
this.refillRate = refillRate;
this.tokens = capacity;
this.lastRefillTime = System.nanoTime();
}
/**
* 尝试获取令牌
* @return 是否获取成功
*/
public synchronized boolean tryAcquire() {
refill();
if (tokens >= 1) {
tokens -= 1;
return true;
}
return false;
}
/**
* 尝试获取多个令牌
*/
public synchronized boolean tryAcquire(int permits) {
if (permits <= 0) {
return true;
}
refill();
if (tokens >= permits) {
tokens -= permits;
return true;
}
return false;
}
/**
* 补充令牌
*/
private void refill() {
long now = System.nanoTime();
long elapsedTime = now - lastRefillTime;
// 计算应补充的令牌数量
double newTokens = (elapsedTime / 1_000_000_000.0) * refillRate;
if (newTokens > 0) {
tokens = Math.min(capacity, tokens + newTokens);
lastRefillTime = now;
}
}
/**
* 获取当前可用令牌数
*/
public synchronized double getAvailableTokens() {
refill();
return tokens;
}
}
使用Guava的RateLimiter实现
import com.google.common.util.concurrent.RateLimiter;
/**
* 使用Guava实现的令牌桶
*/
public class GuavaTokenBucket {
private final RateLimiter rateLimiter;
/**
* @param permitsPerSecond 每秒钟产生的令牌数
*/
public GuavaTokenBucket(double permitsPerSecond) {
this.rateLimiter = RateLimiter.create(permitsPerSecond);
}
/**
* 非阻塞获取令牌
*/
public boolean tryAcquire() {
return rateLimiter.tryAcquire();
}
/**
* 非阻塞获取指定数量的令牌
*/
public boolean tryAcquire(int permits) {
return rateLimiter.tryAcquire(permits);
}
/**
* 阻塞获取令牌（可设置超时时间）
*/
public boolean tryAcquireWithTimeout(int permits, long timeout, TimeUnit unit) {
return rateLimiter.tryAcquire(permits, timeout, unit);
}
/**
* 设置新的速率
*/
public void setRate(double permitsPerSecond) {
rateLimiter.setRate(permitsPerSecond);
}
}
高级版令牌桶（支持预热）
import java.util.concurrent.atomic.AtomicLong;
/**
* 高级令牌桶 - 支持预热和动态调整
*/
public class AdvancedTokenBucket {
// 桶容量
private final long capacity;
// 当前令牌数
private final AtomicLong tokens;
// 最后补充时间
private volatile long lastRefillTime;
// 令牌生成速率（每秒）
private volatile double refillRate;
// 预热时间（毫秒）
private final long warmupPeriod;
// 预热期间的速率（可降低速率）
private volatile boolean warmingUp = true;
// 预热开始时间
private long warmupStartTime;
public AdvancedTokenBucket(long capacity, double refillRate, long warmupPeriodMs) {
this.capacity = capacity;
this.refillRate = refillRate;
this.warmupPeriod = warmupPeriodMs;
this.tokens = new AtomicLong(capacity);
this.lastRefillTime = System.nanoTime();
this.warmupStartTime = System.currentTimeMillis();
}
/**
* 尝试获取令牌（考虑预热）
*/
public boolean tryAcquire() {
refillTokens();
// 预热期间降低获取成功率
if (warmingUp && isInWarmupPeriod()) {
double warmupFactor = getWarmupFactor();
if (Math.random() > warmupFactor) {
return false;
}
}
while (true) {
long currentTokens = tokens.get();
if (currentTokens <= 0) {
return false;
}
if (tokens.compareAndSet(currentTokens, currentTokens - 1)) {
return true;
}
}
}
/**
* 补充令牌
*/
private void refillTokens() {
long now = System.nanoTime();
long elapsedNanos = now - lastRefillTime;
double elapsedSeconds = elapsedNanos / 1_000_000_000.0;
long newTokens = (long) (elapsedSeconds * refillRate);
if (newTokens > 0) {
tokens.updateAndGet(current ->
Math.min(capacity, current + newTokens)
);
lastRefillTime = now;
}
}
/**
* 检查是否在预热期
*/
private boolean isInWarmupPeriod() {
return (System.currentTimeMillis() - warmupStartTime) < warmupPeriod;
}
/**
* 获取预热因子
*/
private double getWarmupFactor() {
long elapsed = System.currentTimeMillis() - warmupStartTime;
// 预热从10%开始，线性增长到100%
return Math.max(0.1, (double) elapsed / warmupPeriod);
}
/**
* 完成预热
*/
public void finishWarmup() {
warmingUp = false;
}
/**
* 获取当前令牌数
*/
public long getAvailableTokens() {
refillTokens();
return tokens.get();
}
}
分布式令牌桶（基于Redis）
import redis.clients.jedis.Jedis;
import redis.clients.jedis.Transaction;
import java.util.List;
/**
* 基于Redis的分布式令牌桶
*/
public class DistributedTokenBucket {
private final Jedis jedis;
private final String bucketKey;
private final String tokenKey;
private final String timestampKey;
private final long capacity;
private final double refillRate;
public DistributedTokenBucket(Jedis jedis, String bucketName,
long capacity, double refillRate) {
this.jedis = jedis;
this.bucketKey = "bucket:" + bucketName;
this.tokenKey = bucketKey + ":tokens";
this.timestampKey = bucketKey + ":timestamp";
this.capacity = capacity;
this.refillRate = refillRate;
// 初始化
if (!jedis.exists(tokenKey)) {
jedis.set(tokenKey, String.valueOf(capacity));
jedis.set(timestampKey, String.valueOf(System.currentTimeMillis()));
}
}
/**
* 使用Lua脚本实现原子操作
*/
public boolean tryAcquire() {
String luaScript =
"local tokens = tonumber(redis.call('get', KEYS[1])) " +
"local lastRefill = tonumber(redis.call('get', KEYS[2])) " +
"local now = tonumber(ARGV[1]) " +
"local rate = tonumber(ARGV[2]) " +
"local capacity = tonumber(ARGV[3]) " +
"local refill = (now - lastRefill) / 1000 * rate " +
"tokens = math.min(capacity, tokens + refill) " +
"if tokens >= 1 then " +
" tokens = tokens - 1 " +
" redis.call('set', KEYS[1], tokens) " +
" redis.call('set', KEYS[2], now) " +
" return 1 " +
"else " +
" redis.call('set', KEYS[1], tokens) " +
" redis.call('set', KEYS[2], now) " +
" return 0 " +
"end";
Object result = jedis.eval(luaScript, 2, tokenKey, timestampKey,
String.valueOf(System.currentTimeMillis()),
String.valueOf(refillRate),
String.valueOf(capacity));
return (Long) result == 1;
}
}
使用示例
public class TokenBucketDemo {
public static void main(String[] args) throws InterruptedException {
// 基础令牌桶示例
System.out.println("=== 基础令牌桶示例 ===");
TokenBucket bucket = new TokenBucket(10, 2); // 容量10，每秒产生2个令牌
// 模拟请求
for (int i = 0; i < 20; i++) {
boolean allowed = bucket.tryAcquire();
System.out.printf("请求%d: %s%n", i, allowed ? "通过" : "被限流");
Thread.sleep(100);
}
// Guava RateLimiter示例
System.out.println("\n=== Guava RateLimiter示例 ===");
GuavaTokenBucket guavaBucket = new GuavaTokenBucket(5); // 每秒5个令牌
for (int i = 0; i < 10; i++) {
boolean allowed = guavaBucket.tryAcquire();
System.out.printf("请求%d: %s%n", i, allowed ? "通过" : "被限流");
Thread.sleep(200);
}
// 高级令牌桶示例（带预热）
System.out.println("\n=== 高级令牌桶示例（预热） ===");
AdvancedTokenBucket advancedBucket = new AdvancedTokenBucket(10, 2, 5000);
for (int i = 0; i < 30; i++) {
boolean allowed = advancedBucket.tryAcquire();
System.out.printf("请求%d: %s (可用令牌: %d)%n",
i, allowed ? "通过" : "被限流",
advancedBucket.getAvailableTokens());
Thread.sleep(200);
}
}
}
实际应用场景
/**
* 实际应用 - API限流
*/
public class APIRateLimiter {
private final TokenBucket tokenBucket;
public APIRateLimiter(int maxRequests, double refillRate) {
this.tokenBucket = new TokenBucket(maxRequests, refillRate);
}
/**
* 处理API请求
*/
public Response processRequest(Request request) {
if (!tokenBucket.tryAcquire()) {
return Response.rateLimited("请求过于频繁，请稍后再试");
}
// 执行业务逻辑
return executeBusinessLogic(request);
}
private Response executeBusinessLogic(Request request) {
// 业务逻辑实现
return Response.success("处理成功");
}
}
/**
* 实际应用 - 用户登录限制
*/
public class LoginRateLimiter {
private final TokenBucket loginBucket;
private final int maxAttempts = 5;
public LoginRateLimiter() {
// 5次尝试，每分钟补充1次
this.loginBucket = new TokenBucket(maxAttempts, 1.0/60);
}
public LoginResult attemptLogin(String username, String password) {
if (!loginBucket.tryAcquire()) {
return LoginResult.tooManyAttempts("登录尝试次数过多，请稍后再试");
}
// 验证用户
return verifyUser(username, password);
}
private LoginResult verifyUser(String username, String password) {
// 用户验证逻辑
return LoginResult.success();
}
}
这个实现涵盖了：
- 基础令牌桶：简单的同步实现
- Guava版本：使用成熟库的实现
- 高级版本：支持预热和动态调整
- 分布式版本：基于Redis支持分布式环境
- 实际应用：API限流和登录限制示例
可以根据具体需求选择合适的实现方式。
