令牌桶(TokenBucket)限流 - Java实现

主要问题: 如何实现以固定的速率向桶中添加令牌?

我们当然不用起一个定时任务不停的向桶中添加令牌,只要在访问时记下访问时间,下次访问时计算出两次访问时间的间隔,然后向桶中补充令牌,补充的令牌数为

(long) (durationMs * rate * 1.0 / 1000)

其中rate是每秒补充令牌数。

以下是完整的实现代码

public class TokenBucketRateLimiter {

    long capacity;                      // 桶的容量
    long rate;                          // 令牌发放速率, permits-per-second
    long currentTokenNum;               // 当前桶中的令牌数量
    long lastAddTokenTime;              // 上次补充令牌的时间

    public TokenBucketRateLimiter(long capacity, long rate) {
        this.capacity = capacity;
        this.rate = rate;
        currentTokenNum = capacity;
        lastAddTokenTime = System.currentTimeMillis();
    }

    public boolean acquire() {
        return acquire(1);
    }

    public boolean acquire(int permits) {
        if (permits > currentTokenNum) {
            long accessTime = System.currentTimeMillis();
            long durationMs = accessTime - lastAddTokenTime;
            long newTokenNum = (long) (durationMs * rate * 1.0 / 1000);
            if (newTokenNum > 0) {
                currentTokenNum = Math.min(currentTokenNum + newTokenNum, capacity);
                this.lastAddTokenTime = accessTime;
            }
            if (permits > currentTokenNum) return false;
        }
        this.currentTokenNum -= permits;
        return true;
    }

}

发布于 04-04