线程安全令牌桶算法实现疑问：多线程并发安全问题探讨

原单节点令牌桶限流实现存在并发安全问题：当tokens值为1时，两个线程可能同时通过tokens.get() > 0的检查，随后依次执行decrementAndGet()导致tokens变为负数，违规允许第二次请求通过。请问是否只能使用synchronized或ReentrantLock API来解决该并发安全问题？

原实现代码 #

public class TokenBucketRateLimiter implements RateLimiter {
    // 桶可容纳的最大令牌数
    private final long threshold;
    // 令牌补充的时间单位（毫秒）
    private final long windowUnit = 1000;
    // 当前可用令牌数
    private final AtomicLong tokens = new AtomicLong(0);
    // 上次补充令牌的时间戳
    private volatile long lastRefillTimestamp = System.currentTimeMillis();

    /**
     * 构造指定最大令牌数的令牌桶限流器
     *
     * @param threshold 桶可容纳的最大令牌数
     */
    public TokenBucketRateLimiter(long threshold) {
        this.threshold = threshold;
        // 初始化时将桶填满
        this.tokens.set(threshold);
    }

    /**
     * 尝试获取一个令牌
     *
     * @return 获取成功返回true，否则返回false
     */
    @Override
    public boolean tryAcquire() {
        long currentTime = System.currentTimeMillis();
        // 计算自上次补充以来应添加的令牌数
        long tokensToAdd = ((currentTime - lastRefillTimestamp) / windowUnit) * threshold;

        // 补充令牌
        if (tokensToAdd > 0) {
            // 确保令牌数不超过桶的最大容量
            tokens.set(Math.min(threshold, tokens.addAndGet(tokensToAdd)));
            // 更新补充时间戳
            lastRefillTimestamp = currentTime;
        }

        // 尝试获取令牌
        if (tokens.get() > 0) {
            // 减少令牌数并允许访问
            tokens.decrementAndGet();
            return true;
        }

        // 令牌桶为空，拒绝访问
        return false;
    }
}

问题解答 #

当然不是只能用synchronized或ReentrantLock，还可以利用AtomicLong的**原子比较并交换（CAS）**操作实现无锁的并发安全控制，这种方式在高并发场景下通常比锁的性能更优。

核心思路是将"检查令牌数是否大于0"和"递减令牌数"这两个操作合并为一个原子操作，避免线程间竞争导致的不一致。以下是两种可行的修改方案：

方案一：使用CAS循环实现 #

@Override
public boolean tryAcquire() {
    long currentTime = System.currentTimeMillis();
    long tokensToAdd = ((currentTime - lastRefillTimestamp) / windowUnit) * threshold;

    if (tokensToAdd > 0) {
        // 先计算新令牌数再原子设置，避免多线程同时累加的问题
        long newTokenCount = Math.min(threshold, tokens.get() + tokensToAdd);
        tokens.set(newTokenCount);
        lastRefillTimestamp = currentTime;
    }

    // 原子尝试获取令牌：只有当前令牌数>0时才递减
    long currentTokens;
    do {
        currentTokens = tokens.get();
        if (currentTokens <= 0) {
            return false;
        }
        // CAS操作：如果当前值仍为currentTokens，就更新为currentTokens-1
    } while (!tokens.compareAndSet(currentTokens, currentTokens - 1));

    return true;
}

方案二：使用getAndUpdate简化逻辑 #

@Override
public boolean tryAcquire() {
    long currentTime = System.currentTimeMillis();
    long tokensToAdd = ((currentTime - lastRefillTimestamp) / windowUnit) * threshold;

    if (tokensToAdd > 0) {
        long newTokenCount = Math.min(threshold, tokens.get() + tokensToAdd);
        tokens.set(newTokenCount);
        lastRefillTimestamp = currentTime;
    }

    // 原子更新令牌数，返回更新前的数值
    long remaining = tokens.getAndUpdate(current -> current > 0 ? current - 1 : current);
    return remaining > 0;
}

关键说明 #

• 两种方案都通过原子操作消除了"检查"和"递减"之间的竞态条件，确保同一时间只有一个线程能成功获取令牌。
• 令牌补充部分也做了优化：先计算最终令牌数再原子设置，避免原代码中addAndGet可能导致的多线程重复累加问题。

内容的提问来源于stack exchange，提问作者curiousengineer