【限流算法】java实现令牌桶算法

本文实现了一种基本的令牌桶算法

令牌桶算法思想：以固定速率产生令牌，放入令牌桶，每次用户请求都得申请令牌，令牌不足则拒绝请求或等待。

代码逻辑：线程池每0.5s发送随机数量的请求，每次请求计算当前的令牌数量，请求令牌数量超出当前令牌数量，则产生限流。

@Slf4j
public class TokensLimiter {

    private ScheduledExecutorService scheduledExecutorService = Executors.newScheduledThreadPool(5);

    // 最后一次令牌发放时间
    public long timeStamp = System.currentTimeMillis();
    // 桶的容量
    public int capacity = 10;
    // 令牌生成速度10/s
    public int rate = 10;
    // 当前令牌数量
    public int tokens;

    public void acquire() {
        scheduledExecutorService.scheduleWithFixedDelay(() -> {
            long now = System.currentTimeMillis();
            // 当前令牌数
            tokens = Math.min(capacity, (int) (tokens + (now - timeStamp) * rate / 1000));
            //每隔0.5秒发送随机数量的请求
            int permits = (int) (Math.random() * 9) + 1;
            log.info("请求令牌数：" + permits + "，当前令牌数：" + tokens);
            timeStamp = now;
            if (tokens < permits) {
                // 若不到令牌,则拒绝
                log.info("限流了");
            } else {
                // 还有令牌，领取令牌
                tokens -= permits;
                log.info("剩余令牌=" + tokens);
            }
        }, 1000, 500, TimeUnit.MILLISECONDS);
    }

    public static void main(String[] args) {
        TokensLimiter tokensLimiter = new TokensLimiter();
        tokensLimiter.acquire();
    }

}

输出结果：

16:13:20.042 [pool-1-thread-1] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：1，当前令牌数：10
16:13:20.045 [pool-1-thread-1] INFO com.example.demo.limit.TokensLimiter - 剩余令牌=9
16:13:20.549 [pool-1-thread-1] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：7，当前令牌数：10
16:13:20.549 [pool-1-thread-1] INFO com.example.demo.limit.TokensLimiter - 剩余令牌=3
16:13:21.054 [pool-1-thread-2] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：5，当前令牌数：8
16:13:21.054 [pool-1-thread-2] INFO com.example.demo.limit.TokensLimiter - 剩余令牌=3
16:13:21.559 [pool-1-thread-1] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：1，当前令牌数：8
16:13:21.559 [pool-1-thread-1] INFO com.example.demo.limit.TokensLimiter - 剩余令牌=7
16:13:22.063 [pool-1-thread-3] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：7，当前令牌数：10
16:13:22.063 [pool-1-thread-3] INFO com.example.demo.limit.TokensLimiter - 剩余令牌=3
16:13:22.568 [pool-1-thread-3] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：7，当前令牌数：8
16:13:22.568 [pool-1-thread-3] INFO com.example.demo.limit.TokensLimiter - 剩余令牌=1
16:13:23.072 [pool-1-thread-3] INFO com.example.demo.limit.TokensLimiter - 请求令牌数：7，当前令牌数：6
16:13:23.072 [pool-1-thread-3] INFO com.example.demo.limit.TokensLimiter - 限流了