令牌桶在Java中的实现

最新推荐文章于 2025-07-02 20:57:20 发布

B1azar

最新推荐文章于 2025-07-02 20:57:20 发布

阅读量1k

点赞数

CC 4.0 BY-SA版权

文章标签： java 开发语言后端大数据

本文链接：https://blog.youkuaiyun.com/B1azar/article/details/126679970

本文介绍了在高并发场景下使用限流算法的重要性，并详细探讨了令牌桶算法的原理和应用。通过分析难点，如如何按给定流速添加令牌和实现线程安全的数据结构，文章提供了一个详细的Java实现令牌桶算法的代码示例。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

限流算法

在高并发场景下，除了使用消息队列、缓存来处理外，我们还可以限定请求的次数，即让我们规定数量的请求进来，于是便有了限流算法，限流可以常用的有：

计数限流
窗口限流
令牌桶
漏桶

令牌桶原理

令牌桶在实际场景中应用更广泛，guava也提供了现成的方法供我们使用，可学习过程中，难免重复造轮子

令牌桶整体的处理如下图所示

在这里插入图片描述

过程如下

初始桶容量，按一定流速向桶中添加令牌
用户发起请求时尝试获取令牌
获取到令牌则进入下一步业务逻辑
如果没有令牌，则视为请求失败

难点

实现过程主要有两个难点

如何按给定流速添加令牌

使用newScheduledThreadPool线程池可以定时执行任务，其中scheduleAtFixedRate方法可以按给定的间隔定期执行相同的任务
桶使用什么数据结构实现

因为对于高并发场景，要保证桶的线程安全，这里可以直接使用阻塞队列实现，如果不想用阻塞队列，也可以自己尝试实现，保证poll和offer操作就可以

详细代码

public class TokenBucket {
    /**
     * capacity: 桶的容量
     * period: 添加令牌间隔
     * amount: 每次添加数量
     * bucket: 桶
     * period/amount: 流速
     */
    private int capacity;
    private int period;
    private int amount;

    BlockingQueue<String> bucket;

    public TokenBucket(int capacity, int period, int amount) {
        this.period = period;
        this.amount = amount;
        this.capacity = capacity;
    }

    /**
     * 初始化令牌桶，手动进行初始化，否则在第一次tryAcquire操作时初始化
     */
    public void init() {
        bucket = new ArrayBlockingQueue<>(capacity);
        scheduledAddToken();
    }

    /**
     * 获取令牌
     * @return
     * 返回true表示获取成功，false表示失败
     */
    public boolean tryAcquire() {
        if (bucket == null) init();
        return bucket.poll() != null;
    }

    /**
     * 添加令牌
     */
    private void addToken() {
        for (int i = 0; i < amount; i++) {
            bucket.offer("token");
        }
    }

    /**
     * 定时进行添加令牌
     */
    private void scheduledAddToken() {
        Executors.newScheduledThreadPool(1)
                .scheduleAtFixedRate(this::addToken, 0, period, TimeUnit.SECONDS);
    }

    public static void main(String[] args) throws InterruptedException {
        int period = 10, capacity = 10, amount = 5;
        TokenBucket bucket = new TokenBucket(capacity, period, amount);
        bucket.init();
        for (int i = 0; i < 100; i++) {
            String s = "第" + i + "次";
            if (bucket.tryAcquire()) {
                System.out.println(s + "拿取成功");
            } else {
                System.out.println(s + "拿取失败");
            }
            Thread.sleep(1000);
        }
    }
}