每日算法 — 使用java实现令牌桶：限流算法与流量整形

引言

在高并发系统中，限流是保护服务稳定性的核心手段。无论是API接口防刷、消息队列削峰，还是资源池访问控制，都需要精确的流量调控机制。令牌桶（Token Bucket）算法以其平滑输出、允许突发流量的特性，成为业界最广泛采用的限流方案。本文用Java实现线程安全的令牌桶，并对比漏桶算法。

令牌桶原理

想象一个以固定速率放入令牌的桶，请求到达时需要获取令牌才能执行。若桶空则请求等待或被拒绝。

核心参数

• 容量（Capacity）：桶中最多存放的令牌数，决定最大突发流量
• 速率（Rate）：每秒产生的令牌数，决定长期平均流量
• 当前令牌数：动态变化，随时间增长、随消费减少

与漏桶算法对比

Java实现

import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

/**
 * 线程安全令牌桶限流器
 */
class TokenBucket {
    private final long capacity;           // 桶容量
    private final long tokensPerSecond;    // 令牌生成速率
    private double availableTokens;        // 当前可用令牌数（double精度避免累积误差）
    private long lastRefillTimestamp;      // 上次补充令牌的时间戳（纳秒）
    private final Lock lock;

    public TokenBucket(long capacity, long tokensPerSecond) {
        this.capacity = capacity;
        this.tokensPerSecond = tokensPerSecond;
        this.availableTokens = capacity;   // 初始满桶
        this.lastRefillTimestamp = System.nanoTime();
        this.lock = new ReentrantLock();
    }

    /**
     * 尝试获取指定数量的令牌
     * @param tokensNeeded 需要的令牌数
     * @return true=获取成功，false=令牌不足
     */
    public boolean tryAcquire(long tokensNeeded) {
        if (tokensNeeded <= 0) return true;
        if (tokensNeeded > capacity) return false; // 请求超过桶容量，直接拒绝

        lock.lock();
        try {
            refill(); // 先补充令牌

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

    /**
     * 获取令牌，若不足则阻塞等待
     */
    public void acquire(long tokensNeeded) throws InterruptedException {
        if (tokensNeeded <= 0) return;

        while (true) {
            lock.lock();
            try {
                refill();

                if (availableTokens >= tokensNeeded) {
                    availableTokens -= tokensNeeded;
                    return;
                }

                // 计算需要等待的时间
                double tokensToWait = tokensNeeded - availableTokens;
                long waitNanos = (long) (tokensToWait * 1_000_000_000L / tokensPerSecond);
                lock.unlock();

                if (waitNanos > 0) {
                    Thread.sleep(waitNanos / 1_000_000, (int) (waitNanos % 1_000_000));
                }
            } catch (Exception e) {
                if (lock.tryLock()) lock.unlock();
                throw e;
            }
        }
    }

    /**
     * 根据时间流逝补充令牌
     */
    private void refill() {
        long now = System.nanoTime();
        long elapsedNanos = now - lastRefillTimestamp;

        if (elapsedNanos > 0) {
            // 计算这段时间应生成的令牌数
            double tokensToAdd = elapsedNanos * tokensPerSecond / 1_000_000_000.0;
            availableTokens = Math.min(capacity, availableTokens + tokensToAdd);
            lastRefillTimestamp = now;
        }
    }

    public double getAvailableTokens() {
        lock.lock();
        try {
            refill();
            return availableTokens;
        } finally {
            lock.unlock();
        }
    }
}

/**
 * 无锁版令牌桶：基于AtomicLong的时间戳实现（单线程场景优化）
 */
class LockFreeTokenBucket {
    private final long capacity;
    private final long tokensPerSecond;
    private final AtomicLong state; // 高32位存储lastRefill的秒级时间，低32位存储令牌数*1000

    public LockFreeTokenBucket(long capacity, long tokensPerSecond) {
        this.capacity = capacity;
        this.tokensPerSecond = tokensPerSecond;
        long initialTokens = capacity * 1000;
        long initialTime = System.currentTimeMillis() << 32;
        this.state = new AtomicLong(initialTime | initialTokens);
    }

    public boolean tryAcquire(long tokensNeeded) {
        if (tokensNeeded > capacity) return false;

        while (true) {
            long current = state.get();
            long lastTime = current >>> 32;
            long currentTokensMilli = current & 0xFFFFFFFFL;
            double currentTokens = currentTokensMilli / 1000.0;

            long now = System.currentTimeMillis();
            long elapsed = now - lastTime;
            double newTokens = Math.min(capacity,
                currentTokens + elapsed * tokensPerSecond / 1000.0);

            if (newTokens < tokensNeeded) return false;

            long newTokensMilli = (long) ((newTokens - tokensNeeded) * 1000);
            long newState = (now << 32) | (newTokensMilli & 0xFFFFFFFFL);

            if (state.compareAndSet(current, newState)) {
                return true;
            }
        }
    }
}

/**
 * 多线程限流测试
 */
class RateLimiterTest {
    private final TokenBucket bucket;
    private final AtomicLong successCount = new AtomicLong(0);
    private final AtomicLong failCount = new AtomicLong(0);

    public RateLimiterTest(long capacity, long rate) {
        this.bucket = new TokenBucket(capacity, rate);
    }

    public void runTest(int threadCount, int requestsPerThread) throws InterruptedException {
        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
        CountDownLatch latch = new CountDownLatch(threadCount);

        long startTime = System.currentTimeMillis();

        for (int i = 0; i < threadCount; i++) {
            executor.submit(() -> {
                for (int j = 0; j < requestsPerThread; j++) {
                    if (bucket.tryAcquire(1)) {
                        successCount.incrementAndGet();
                    } else {
                        failCount.incrementAndGet();
                    }
                    // 模拟请求间隔
                    try { Thread.sleep(10); } catch (InterruptedException e) { break; }
                }
                latch.countDown();
            });
        }

        latch.await();
        executor.shutdown();

        long duration = System.currentTimeMillis() - startTime;
        System.out.printf("测试完成：%d线程 × %d请求 = %d总请求%n",
            threadCount, requestsPerThread, threadCount * requestsPerThread);
        System.out.printf("成功: %d | 拒绝: %d | 耗时: %dms%n",
            successCount.get(), failCount.get(), duration);
        System.out.printf("实际通过率: %.2f%%%n",
            100.0 * successCount.get() / (successCount.get() + failCount.get()));
    }
}

/**
 * 测试入口
 */
public class TokenBucketDemo {
    public static void main(String[] args) throws InterruptedException {
        System.out.println("=== 基础令牌桶测试 ===");
        TokenBucket bucket = new TokenBucket(10, 5); // 容量10，每秒5个

        // 初始 burst：桶满，前10个请求应全部通过
        System.out.println("初始令牌数: " + bucket.getAvailableTokens());
        for (int i = 1; i <= 15; i++) {
            boolean ok = bucket.tryAcquire(1);
            System.out.printf("请求%02d: %s (剩余: %.1f)%n", i,
                ok ? "通过" : "拒绝", bucket.getAvailableTokens());
        }

        // 等待补充
        System.out.println("\n等待2秒补充令牌...");
        Thread.sleep(2000);
        System.out.println("2秒后令牌数: " + bucket.getAvailableTokens());

        System.out.println("\n=== 并发限流测试 ===");
        RateLimiterTest test = new RateLimiterTest(20, 10); // 容量20，每秒10个
        test.runTest(5, 20); // 5线程，每线程20请求
    }
}

算法要点解析

时间精度处理

令牌桶的关键在于精确计算时间差对应的令牌增量。使用System.nanoTime()而非currentTimeMillis()，因为：
– nanoTime()单调递增，不受系统时间调整影响
– 纳秒级精度避免高速率场景下的累积误差

突发流量控制

容量参数决定系统允许的突发上限。例如容量100、速率10/秒的系统：
– 空闲一段时间后桶满（100令牌）
– 突发100个请求可立即处理
– 之后被限制为10/秒

预热与冷启动

生产环境中常需要预热模式：系统刚启动时限流更严格，逐步放开。可通过动态调整速率或初始令牌数实现。

复杂度分析

总结

令牌桶算法以简洁的数学模型实现了灵活的流量控制。在实际生产环境中，Guava的RateLimiter就是基于此原理。理解其实现有助于根据业务场景定制限流策略，如区分用户等级的差异化限速、接口级别的细粒度控制等。