Java漏桶算法的简单代码实例

@FunctionalInterface
public interface Callback<Task> {
    void process(Task task) throws InterruptedException;
}

/**
 * <p>简单漏桶算法实现</p>
 */
public class LeakyBucketHandler<T> {
	// 漏水速率 /s（TPS/QPS）
	private Integer rate;
	private long lastTime = System.currentTimeMillis();
	private final int capacity;
	// 最大并发量（桶最大容量）
	private final ArrayBlockingQueue<T> queue;
	// init
	public LeakyBucketHandler(Integer rate, int capacity) {
		this.rate = rate;
		this.capacity = capacity;
		this.queue = new ArrayBlockingQueue<T>(capacity);
	}
	public LeakyBucketHandler(int capacity) {
		this.capacity = capacity;
		this.queue = new ArrayBlockingQueue<T>(capacity);
	}
	public boolean acquire(T b) {
		if (queue.size() > capacity) {
			return false;
		} else {
			queue.offer(b);
			return true;
		}
	}
	public synchronized void consume(Callback<T> callBack) throws InterruptedException {
		if (rate == null || rate < 1) {
			throw new IllegalArgumentException("rate value is" + rate);
		}
		while (queue.size() > 0) {
			long now = System.currentTimeMillis();
			if ((now-lastTime) > 1000/rate) {
				T t = queue.poll();
				System.out.println("interval-" + (now - lastTime + "-ms"));;
				lastTime = now;
				callBack.process(t);
				System.out.println("bucket size is " + queue.size());
			} else {
				Thread.sleep(1);
			}
		}
	}
	public Integer getQueueSize(){
		return queue.size();
	}
}

public class Demo {
    public static void main(String[] args) throws InterruptedException {
        // 获取任务队列
        List<String> taskList = getTaskList();
        LeakyBucketHandler<String> leakyBucket = new LeakyBucketHandler<>(5, taskList.size());
        for (String s : taskList) {
            leakyBucket.acquire(s);
        }
        System.out.println("leakyBucket.getQueueSize()="+leakyBucket.getQueueSize());
        leakyBucket.consume(task -> {
            CompletableFuture.runAsync(()->{
                System.out.println("消费桶中对象---"+task+"开始");
                try {
                    // 模拟业务耗时
                    Thread.sleep(200);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
            });
        });
    }
    // 一般从数据库或其他数据源拉取数据，这里mock一下
    private static List<String> getTaskList() {
        List<String> list = new ArrayList<>();
        for (int i = 0; i < 100; i++) {
            list.add(i + "");
        }
        return list;
    }
}