令牌桶算法简洁(简陋)实现

令牌桶算法简洁(简陋)实现

参考了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;
    }
}

你可能感兴趣的:(令牌桶算法简洁(简陋)实现)