限流 -- 令牌桶算法

一、什么是令牌桶算法？

令牌桶算法是一种流量控制策略，用于限制请求数。

它通过维护一个固定容量的“令牌桶”，并以恒定速率向其中添加令牌。当请求到达时，需要从桶中取出一个令牌才能被处理，如果没有可用的令牌，则请求将被拒绝或等待

二、使用 Java 实现的令牌桶简单示例：

// 声明一个令牌桶
public class TokenBucket {
    private final int capacity; // 仓库容量
    private final AtomicInteger tokens; // 当前剩余token数量
    private final int refillRate; // 每秒补充token数

    public TokenBucket(int capacity, int refillRate) {
        this.capacity = capacity;
        this.tokens = new AtomicInteger(capacity);
        this.refillRate = refillRate;

        ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
        // 每隔 (1000 / refillRate) 毫秒补充一个令牌
        scheduler.scheduleAtFixedRate(() -> {
            if (tokens.get() < capacity) {
                tokens.incrementAndGet();
            }
        }, 0, 1000 / refillRate, TimeUnit.MILLISECONDS);
    }

    // 从令牌桶中消耗一个令牌
    public synchronized boolean tryConsume() {
        if (tokens.get() > 0) {
            tokens.decrementAndGet();
            return true;
        }
        return false;
    }

    public static void main(String[] args) throws InterruptedException {
        TokenBucket tokenBucket = new TokenBucket(5, 2);
        ExecutorService threadPool = Executors.newFixedThreadPool(10);
        ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);

        // 每隔1s钟，再发起三个请求， 按照令牌桶的策略，之后的每秒，每次三个请求，都会被接收2个
        scheduler.scheduleAtFixedRate(() -> {
            CountDownLatch countDownLatch = new CountDownLatch(3);
            System.out.println("--------------------------start-------------------------------------->");
            for (int i = 0; i < 3; i++) {
                threadPool.execute(() -> {
                    System.out.println(new Date() + "Request " + Thread.currentThread().getId() + ": " + (tokenBucket.tryConsume() ? "Accepted" :
                            "Rejected"));
                    countDownLatch.countDown();
                });
            }
            try {
                countDownLatch.await();
            } catch (InterruptedException e) {

            }
            System.out.println("--------------------------end-------------------------------------->");
        }, 0, 1000, TimeUnit.MILLISECONDS);

    }
}

在这个示例中，我们创建了一个容量为5的令牌桶，并设置每秒补充2个令牌。然后定时每秒发起3个请求，观察结果可以看到，前三秒的请求都被接收，后面所有的每秒请求，接收2个，拒绝一个

三、单个Jvm应用的限流框架

在 Java 生态中，有一些成熟的库实现了令牌桶算法，以下是两个常用的库：

1、Guava RateLimiter

Guava RateLimiter是一个用于限制操作速率的工具类，它可以帮助我们在需要控制某些操作频率时提供稳定的请求处理速度

public class RateLimiterDemo {

    public static void main(String[] args) {
        // 创建一个每秒允许2个令牌（即两次操作）通过的RateLimiter实例
        RateLimiter rateLimiter = RateLimiter.create(2.0);

        for (int i = 1; i <= 10; i++) {
            // 请求获取令牌，如果没有可用令牌，则阻塞等待
            if (rateLimiter.tryAcquire()) {

                System.out.println("Task " + i + " acquired");
                // 执行任务逻辑（此处为了演示效果，只打印输出）
                doTask(i);
            } else {
                System.out.println("Task " + i + " reject");
            }
        }
    }

    private static void doTask(int taskId) {
        System.out.println("Executing task: " + taskId);
    }
}

2、Bucket4j

Bucket4j 是另一个专门为 Java 设计的高性能、灵活且易于使用的令牌桶限流库。它支持多种存储后端（如 In-Memory、Hazelcast 和 Ignite），并允许分布式部署。

文档：https://github.com/vladimir-bukhtoyarov/bucket4j/blob/master/doc-pages/basic-usage.md

简单实用：

<dependency>
    <groupId>com.github.vladimir-bukhtoyarov</groupId>
    <artifactId>bucket4j-core</artifactId>
    <version>6.0.2</version>
</dependency>

使用以下代码创建一个基于内存（非分布式）的RateLimiter：

import io.github.bucket4j.*;

public class Bucket4JDemo {

    public static void main(String[] args) {
        // 创建每秒允许2个令牌通过（即两次操作）的RateLimiter实例
        Bucket rateLimiter = createRateLimiter(2);

        for (int i = 1; i <= 10; i++) {
            // 请求获取令牌，如果没有可用令牌，则阻塞等待
            long waitTimeNanos = rateLimiter.asScheduler().consume(1);
            System.out.println("Task " + i + " acquired at nanosecond: " + waitTimeNanos);

            // 执行任务逻辑（此处为了演示效果，只打印输出）
            doTask(i);
        }
    }

    private static Bucket createRateLimiter(int permitsPerSecond) {
        // 定义令牌桶配置
        Bandwidth limit = Bandwidth.simple(permitsPerSecond, Duration.ofSeconds(1));

        // 创建基于内存的Bucket实例
        return Bucket4j.builder().addLimit(limit).build();
    }

    private static void doTask(int taskId) {
      System.out.println("Executing task: " + taskId); 
       System.out.println("Executing task: " + taskId); 
     }
}

在这个示例中，我们创建了一个每秒允许2个令牌通过的Bucket4j RateLimiter实例。然后，在循环中执行10次任务，并在每次执行前请求获取令牌。

由于Bucket4j限制了每秒最多只能有2个令牌通过，因此上述代码会以大约0.5秒一次的速度执行任务。

注意：consume()方法会阻塞等待直到有可用的令牌。如果你不希望阻塞，可以使用tryConsume()方法尝试获取令牌，并在没有可用令牌时立即返回false。

这个示例展示了如何使用Bucket4j进行基本限流操作。对于分布式场景和更高级功能，请参考Bucket4j文档。

四、分布式的限流框架