主要问题: 如何实现以固定的速率向桶中添加令牌?
我们当然不用起一个定时任务不停的向桶中添加令牌,只要在访问时记下访问时间,下次访问时计算出两次访问时间的间隔,然后向桶中补充令牌,补充的令牌数为
durationMs * rate / 1000
以下是完整的实现代码
public class TokenBucketRateLimiter {
long capacity; // 桶的容量
long rate; // 令牌发放速率, permits-per-second
long currentTokenNum; // 当前桶中的令牌数量
long lastAccessTime; // 上次访问令牌桶获取令牌的时间
public TokenBucketRateLimiter(long capacity, long rate) {
this.capacity = capacity;
this.rate = rate;
currentTokenNum = 0;
lastAccessTime = System.currentTimeMillis();
}
public boolean acquire() {
return acquire(1);
}
public boolean acquire(int permits) {
long accessTime = System.currentTimeMillis();
long durationMs = accessTime - lastAccessTime;
currentTokenNum = Math.min(currentTokenNum + durationMs * rate / 1000, capacity);
if (permits > currentTokenNum) return false;
else {
currentTokenNum = currentTokenNum - permits;
this.lastAccessTime = accessTime;
return true;
}
}
}
以上实现可以进一步优化,我们不用每次取令牌时都计算并补充令牌,而是等令牌数不足时,一次性补充令牌。