限流之漏桶算法

java 复制代码

package com.gewb.test;

import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: James.Ge
 * @Date: 2025/12/28
 * @Description 漏桶算法
 */
public class LeakyBucketRateLimiter {
    /**
     * 漏桶容量
     */
    private final int bucketCapacity;

    /**
     * 请求处理速率
     */
    private final int consumeRate;

    /**
     * 请求队列
     */
    private final BlockingQueue<Runnable> bucket;

    private final ScheduledExecutorService executorService;

    public LeakyBucketRateLimiter(int bucketCapacity, int consumeRate) {
        this.bucketCapacity = bucketCapacity;
        this.consumeRate = consumeRate;
        this.bucket = new LinkedBlockingQueue<>(bucketCapacity);
        this.executorService = Executors.newScheduledThreadPool(1);
        // 启动定期消费请求的线程
        startConsuming();
    }

    private void startConsuming() {
//        ScheduledExecutorService executorService = Executors.newScheduledThreadPool(1);
        executorService.scheduleAtFixedRate(() -> {
            for (int i = 0; i < consumeRate; i++) {
                Runnable poll = bucket.poll();
                if(poll != null) {
                    poll.run();
                }
            }
        }, 0, 1, TimeUnit.SECONDS);
    }

    public boolean addRequest(Runnable request) {
        if(bucket.offer(request)) {
            System.out.println("请求已经添加到桶");
            return true;
        } else {
            System.out.println("请求被丢弃，桶已满");
            return false;
        }
    }

    /**
     * 关闭限流器，释放资源
     */
    public void shutdown() {
        while(!bucket.isEmpty()) {
            try {
                Thread.sleep(100);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        executorService.shutdown();

            // 在异常时可保证能够正常释放资源
//            if (!executorService.awaitTermination(5, TimeUnit.SECONDS)) {
//                executorService.shutdownNow();
//            }
//        } catch (InterruptedException e) {
//            executorService.shutdownNow();
//            Thread.currentThread().interrupt();
//        }
        }


        public static void main(String[] args) {
        LeakyBucketRateLimiter rateLimiter = new LeakyBucketRateLimiter(5, 2);
        for (int i = 0; i < 10; i++) {
            int finalI = i;
            rateLimiter.addRequest(() -> {
                LocalDateTime now = LocalDateTime.now();
                String format = now.format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS"));
                System.out.println("处理第" + finalI + "个请求：" + format);
            });
        }

//        try {
//            Thread.sleep(5000);
//        } catch (InterruptedException e) {
//            Thread.currentThread().interrupt();
//        }
        rateLimiter.shutdown();
        System.out.println("程序结束");
    }
}

运行结果：

"C:\Program Files\Java\jdk-17\bin\java.exe" "-javaagent:D:\JetBrains\IntelliJ IDEA 2023.1.2\lib\idea_rt.jar=63321:D:\JetBrains\IntelliJ IDEA 2023.1.2\bin" -Dfile.encoding=UTF-8 -classpath D:\ideaWorkspace2023.1\jdk17\out\production\jdk17;D:\ideaWorkspace2023.1\jdk17\lib\caffeine-3.2.0.jar;D:\ideaWorkspace2023.1\jdk17\lib\jspecify-1.0.0.jar;D:\ideaWorkspace2023.1\jdk17\lib\error_prone_annotations-2.36.0.jar com.gewb.test.LeakyBucketRateLimiter

请求已经添加到桶

请求已经添加到桶

请求已经添加到桶

请求已经添加到桶

请求已经添加到桶

请求已经添加到桶

请求被丢弃，桶已满

请求被丢弃，桶已满

请求被丢弃，桶已满

请求被丢弃，桶已满

处理第0个请求：2025-12-29 11:37:58.659

处理第1个请求：2025-12-29 11:37:58.680

处理第2个请求：2025-12-29 11:37:59.664

处理第3个请求：2025-12-29 11:37:59.664

处理第4个请求：2025-12-29 11:38:00.662

处理第5个请求：2025-12-29 11:38:00.662

程序结束

Process finished with exit code 0