令牌桶限流器学习及分享

分享一个在公司项目中学习到的令牌桶限流器~

文章中的令牌桶限流器主要的应用场景在于告警日志打印、告警、异步任务消费、QPS限制等一些非高并发或需要及时响应的场景。

1.首先列出令牌桶限流器类的核心状态变量：

private final long capacity; // 桶最大容量
private final double refillRate; // 每秒补充的令牌数（速率）
private volatile double currentTokens; // 当前令牌数
private volatile long lastRefillTime; // 上次补充时间（纳秒）

private final ReentrantLock lock = new ReentrantLock();
private final Condition condition = lock.newCondition();

（1）capacity：最多允许多少个请求机会，假设为1000，第 1001 个必须等待。

（2）refillRate：每秒新增多少个令牌进入桶中（即每秒允许多少个请求通过）

（3）currentTokens：当前能被消耗的"通行证"数量，是实际限流控制的核心变量。

（4）lastRefillTime：上一次补充令牌的时间戳（单位：纳秒）。

（5）ReentrantLock：控制对共享变量的并发访问，确保修改/读取是原子的。

（6）condition：在获取不到令牌时挂起线程，等待新令牌补充后被唤醒重试。

2.下面是阻塞获取令牌的代码：

public void acquire() throws InterruptedException {
    lock.lock();
    try {
        while (!tryAcquire(1)) {
            // 计算需要等待的时间（基于令牌补充速率）
            double waitNanos = calculateWaitNanos();
            if (waitNanos <= 0) {
                continue;
            }
            // 限时等待，避免虚假唤醒
            condition.await((long)waitNanos, TimeUnit.NANOSECONDS);
        }
    } finally {
        lock.unlock();
    }
}

说明：线程进来时首先加锁，lock上。这一步是为了保证对令牌数、上次补充时间，这个2共享变量的线程安全操作。如果获取令牌失败的话，会计算等待时间，等待时间大于0时，表示目前时间段令牌已经用完，并通过 Condition.await() 进入限时等待。

3.下面是获取指定数量的令牌的代码：

public boolean tryAcquire(int tokens) {
    lock.lock();
    try {
        refillTokens(); // 先补充令牌
        if (currentTokens >= tokens) {
            currentTokens -= tokens;
            return true;
        }
        return false;
    } finally {
        lock.unlock();
    }
}

说明：首先会进入补充令牌的逻辑，这段逻辑下面再说。if块中如果剩余令牌大于需求令牌数，直接扣减返回true。

4.下面是获取补充令牌的代码：

private void refillTokens() {
    long now = System.nanoTime();
    double elapsed = (now - lastRefillTime) / 1_000_000_000.0; // 转换为秒
    double newTokens = elapsed * refillRate; // 计算应补充的令牌数

    if (newTokens > 0) {
        currentTokens = Math.min(capacity, currentTokens + newTokens);
        lastRefillTime = now;
        condition.signalAll(); // 唤醒等待线程
    }
}

说明：计算上一次补充令牌的时间，计算出需要补充多少令牌。如果需要补充，刷新当前令牌数，然后唤醒等待的线程。

5.最后给出获取令牌的等待时间代码：

private double calculateWaitNanos() {
    double deficit = 1 - currentTokens; // 缺少的令牌数
    return (deficit / refillRate) * 1_000_000_000L; // 转换为纳秒
}

注明：```` if (waitNanos <= 0) { continue; } ``` ` 这段代码是为了防止虚假唤醒的一种防御性编程，我一开始看了半天也没看懂，感觉这个IF永远不会触发啊。后来GPT给我讲明白了。

最后贴一个这个令牌桶限流器的流程图：