令牌桶算法简洁(简陋)实现
参考了guava RateLimiter,实现过程中有几个细节需要注意:
- 发令牌不需要单独起线程,太重了,也太耗费资源了。请求时根据时间计算需要生成多少令牌即可。
- 令牌桶中的令牌不能超过令牌桶的容量。
- 可以通过借用下一个时间片的令牌进行优化,使得令牌借用更加平滑。
class RateLimiter {
// 每秒钟允许的命令数
private long limitPerSec;
// 当前令牌桶中的令牌数
private double currentPermits;
// 上次生成令牌的时间
private long lastAcquireTimestamp;
public RateLimiter(long limitPerSec) {
this.limitPerSec = limitPerSec;
this.lastAcquireTimestamp = System.currentTimeMillis();
}
public synchronized boolean acquire(int permits) {
long now = System.currentTimeMillis();
long duration = now - lastAcquireTimestamp;
// 预先借用了下一个时间片的数据 duration 会小于0
if (duration < 0) {
if (currentPermits >= permits) {
currentPermits -= permits;
return true;
} else {
return false;
}
}
// 增加令牌
lastAcquireTimestamp = now;
currentPermits += duration * (limitPerSec * 1.0 / 1000);
// 令牌桶中的令牌不能超过最大数量
if (currentPermits > limitPerSec) {
currentPermits = limitPerSec;
}
if (currentPermits >= permits) {
currentPermits -= permits;
} else {
// 预先借用下一个时间片的值,一个时间片指生产permits个令牌耗费的时间
lastAcquireTimestamp += 1000 / limitPerSec * permits;
}
return true;
}
}