漏桶算法的Java代码实现

漏桶算法简述

最近要对一个产品进行压力测试，需要生成特定流量的数据，搜索了一下，决定使用漏桶算法。另一种经常使用的令牌桶算法，我认为比较适合流量速率变化较大，希望控制总流量，但对速度波动比较宽容的场合。下面简述一下漏桶算法。
当需要对网络流量进行限制时，即希望在单位时间内只允许流出预设大小的数据量，可以采用漏桶算法。漏桶是个形象且贴切的算法表现。假想有一只下面漏了一个洞的水桶，这个洞漏水的速度忽快忽慢，只要能做到在单位时间内最多只往桶中倒入给定数量的水，无论漏水速度如何变化，在宏观时间范围内漏出的水量都不可能超过这段时间的进水量。但如果水桶容量较大，在微观时间范围内还是有可能出现漏水速度大于进水速度的。
把你想要控制的网络流量当成漏出去的水，把你能控制的一个变量当成桶中的水，这样一来，通过计时加水控制漏水速度的方法，自然就变成了根据时间给变量增加数值来控制网络流量的算法，是不是很容易理解？

实现代码

参考网络资料，稍加修改，形成了如下的算法实现。
漏桶类：

package com.test.test;

/**
 * 流量控制漏桶算法
 * 
 * @author hling
 */
public class Funnel {

	private static final long NANO = 1000000000;

	private long capacity;
	private long leftQuota;
	private long leakingTs;
	private int rate;

	/**
	 * 构造函数
	 * 
	 * @param capatity 容量
	 * @param rate 每秒漏水数量
	 */
	public Funnel(int capatity, int rate) {
		this.capacity = capatity;
		this.leakingTs = System.nanoTime();
		this.rate = rate;
	}

	/**
	 * 补水
	 */
	private void makeSpace() {
		long now = System.nanoTime();
		long time = now - leakingTs;
		long leaked = time * rate / NANO;
		if (leaked < 1) {
			return;
		}
		leftQuota += leaked;

		if (leftQuota > capacity) {
			leftQuota = capacity;
		}
		leakingTs = now;
	}
	
	/**
	 * 漏水。桶里水量不够就返回false
	 * @param quota 漏水量
	 * @return 是否漏水成功
	 */
	public boolean tryWatering(int quota) {
		makeSpace();
		long left = leftQuota - quota;
		if (left >= 0) {
			leftQuota = left;
			return true;
		}
		return false;
	}

	/**
	 * 漏水。没水就阻塞直到蓄满足够的水
	 * @param quota 要漏的数量
	 */
	public void watering(int quota) {
		long left;
		try {
			do {
				makeSpace();
				left = leftQuota - quota;
				if (left >= 0) {
					leftQuota = left;
				} else {
					Thread.sleep(1);
				}
			} while (left < 0);
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
	}

}

测试程序：

package com.test.test;

import java.util.Date;
import java.util.Random;
import java.util.Timer;
import java.util.TimerTask;

public class Sender {
	/** 纳秒 */
	private static final long NANO = 1000000000;

	private volatile static long totalCount = 0;
	private volatile static long totalSize = 0;

	private final static int MAX_PACKET_SIZE_KB = 100; // 每个数据包50KB
	private final static int THRESHOLD_MB = 50; // 每秒50MB流量

	public static final void main(String[] args) {
		startLogger();

		// 测试数据包
		Random rand = new Random();
		// 构造漏桶，这里设置为3个水滴（数据包）大小， 应根据实际情况调整，不能小于1，否则相当于桶容量不到一滴水的量，当然就无法漏水，即流量为0
		// 容量越小，流量在微观时间片内的速率越不会超出流量限制，但在水滴大小变化较大、剩余水量刚好不够一滴水
		// 等情况下，很容易进入等待补水的状态，导致速度低于预设值；
		// 容量越大，宏观速度越稳定接近设定值，但细粒度时间内速度波动范围可能比较大，短时间内漏快了（超速）也可以
		Funnel funnel = new Funnel(3 * MAX_PACKET_SIZE_KB * 1024, THRESHOLD_MB * 1024 * 1024);
		do {
			byte[] data = new byte[rand.nextInt(100) * 1024];
			funnel.watering(data.length);

			totalCount++;
			totalSize += data.length;
		} while (true);
	}

	/**
	 * 数据统计线程
	 */
	private static void startLogger() {
		Timer logTimer = new Timer(true);
		logTimer.schedule(new TimerTask() {
			private long lastTs = 0;
			private long lastCount = 0;
			private long lastSize = 0;

			@Override
			public void run() {
				// 打印周期平均速度
				if (lastTs == 0) {
					// 跳过第一次
					lastTs = System.nanoTime();
					return;
				}

				long tempCount = totalCount;
				long tempSize = totalSize;
				long tempTs = System.nanoTime();

				long count = tempCount - lastCount;
				long size = tempSize - lastSize;
				long duration = tempTs - lastTs;
				float byteSpeed = (float) size * NANO / duration / (1024 * 1024);

				lastCount = tempCount;
				lastSize = tempSize;
				lastTs = tempTs;

				String log = String.format("最近%d秒共产生%d个包，%fMB，平均速度：%dMB/秒", (int) (1.0f * duration / NANO + 0.5), count,
						(float) size / (1024 * 1024), Math.round(byteSpeed));
				System.out.println(log);
			}
		}, new Date(), 5000);
	}

}

程序说明：

• 漏桶类：跟参考资料基本相同，主要是把毫秒计时改成了纳秒计时，因为System.currentTimeMillis()受系统时钟影响，会导致差值计算不稳定，所以改成了System.nanoTime()。按照JDK文档，后者更适合用于计算时间差。此外，漏水函数写了非阻塞和阻塞两种，阻塞方法使用起来更简单。
• 测试程序：做了一个异步线程进行数据统计展示，能直观验证算法的准确性。

因为程序中的加水、漏水都不是原子操作，故程序是不支持多线程的，使用时请注意。如果想要线程安全，最简单的的修改当然就是加锁了，不过可能会影响性能，有需要的朋友自己试试吧。

参考资料

漏桶算法工具类
漏洞算法实现漏斗限流