自定义java令牌桶限流器

public class TokenBucketLimiter {
    private final long rate; // 令牌生成速率（个/秒）
    private final long capacity; // 令牌桶容量
    private long tokens; // 当前令牌数
    private long lastRefillTime; // 上次令牌补充时间
    private final ReentrantLock lock = new ReentrantLock();

    public TokenBucketLimiter(long rate, long capacity) {
        this.rate = rate;
        this.capacity = capacity;
        this.tokens = capacity; // 初始令牌数为桶容量
        this.lastRefillTime = System.nanoTime();
    }

    /**
     * 尝试获取令牌
     *
     * @param tokens 需要的令牌数
     * @return 是否获取成功
     */
    public boolean tryAcquire(int tokens) {
        lock.lock();
        try {
            refill(); // 补充令牌

            if (this.tokens >= tokens) {
                this.tokens -= tokens;
                return true;
            }
            return false;
        } finally {
            lock.unlock();
        }
    }

    /**
     * 补充令牌（根据时间差计算应生成的令牌数）
     */
    private void refill() {
        long now = System.nanoTime();
        long elapsedNanos = now - lastRefillTime;

        // 计算这段时间应生成的令牌数
        long newTokens = (elapsedNanos / 1_000_000_000) * rate;
        if (newTokens > 0) {
            // 令牌数不超过桶容量
            tokens = Math.min(capacity, tokens + newTokens);
            lastRefillTime = now;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        // 每秒生成2个令牌，桶容量为3
        TokenBucketLimiter limiter = new TokenBucketLimiter(2, 3);

        // 模拟请求
        for (int i = 0; i < 50; i++) {
            boolean acquired = limiter.tryAcquire(1);
            System.out.println(String.format("第%d次请求：%s", i + 1, acquired ? "通过" : "被限流"));
            TimeUnit.MILLISECONDS.sleep(250); // 模拟请求间隔
        }
    }
}