网络下载限流之令牌桶算法

什么是限流

一个典型的例子是，当我们在使用百度网盘下载某个文件，在没有开通VIP的情况下，下载速度只有几百K甚至几十K每秒；而我们的网速现在大多数都是10~100M，下载速度和网络出现很大的差距。在开通VIP后下载，速度就会飙升到几M每秒，像这种情况，一般都是软件的客户端或者服务端使用了限流，不花钱就让你慢慢等待。

如何在我们的软件中做限流

令牌桶算法

先上代码，实现一个简易的令牌桶限流器。

//RateLimiter.h
class RateLimiter
{
public:
	RateLimiter();
	//设置每秒可获取令牌数
	void set_rate(double permits_per_second);
	double get_rate();
	//获取一个
	__int64 acquire();
	//获取多个
	__int64 acquire(int permits);
	bool try_acquire(int timeout);
	bool try_acquire(int permis,int timeout);

private:

	//获取需要等待的时间，并更新桶模拟计数
	__int64 claim_next(double permits);
	//同步更新存储的令牌数
	void resync(__int64 now);
	//生成一张令牌时间 ms
	double m_build_interval_ms;
	//下一次可直接获取令牌的时间
	__int64 m_next_free_ms;
	//桶内存储的令牌数
	double m_stored_permits;
	//最大令牌数
	double m_max_permits;
	bool mb_unlimited;
	HMutex m_mutex;
};

//RateLimiter.cpp
#define MIN(a,b)    (((a) < (b)) ? (a) : (b))

RateLimiter::RateLimiter()
{
	m_build_interval_ms = 0.0;
	m_stored_permits = 0.0;
	m_next_free_ms = 0;
	mb_unlimited = false;
}

void RateLimiter::set_rate(double permits_per_second)
{
	if (permits_per_second <= 0.0)
	{
		mb_unlimited = true;
		m_max_permits = permits_per_second;
		return ;
	}
	if (permits_per_second == m_max_permits)
	{
		return ;
	}
	
	mb_unlimited = false;
	double old_max_permits = m_max_permits;
	m_max_permits = permits_per_second;
	m_stored_permits = (old_max_permits == 0.0)? 0.0 
		: m_stored_permits * m_max_permits/old_max_permits;  //如果没有设置则存储数量为0，如果已经设置过，则根据比例更改未消耗完的令牌数量
	
	m_build_interval_ms = 1000.0/permits_per_second;  //每一份需要多少毫秒
}

double RateLimiter::get_rate()
{
	if (mb_unlimited)
	{
		return 0.0;
	}
	return 1000.0/m_build_interval_ms;
}


__int64 RateLimiter::acquire()
{
	return acquire(1);
}

__int64 RateLimiter::acquire(int permits)
{
	if (mb_unlimited)
	{
		return 0;
	}
	
	__int64 i_wait = claim_next(permits);
	if (i_wait > 0)
	{
		HEnvironment::Sleep((unsigned int)i_wait);
	}

	return i_wait;
	
}

bool RateLimiter::try_acquire(int timeout)
{
	return try_acquire(1,timeout);
}

bool RateLimiter::try_acquire(int permis,int timeout)
{
	__int64 now = HEnvironment::get_tick_count64();
	if (m_next_free_ms > now + timeout)
	{
		return false;
	}
	else
	{
		acquire(permis);
	}

	return true;
}
__int64 RateLimiter::claim_next(double permits)
{
	HAutoMutex am(&m_mutex);
	__int64 now = HEnvironment::get_tick_count64();

	resync(now);
	__int64 i_wait_ms = m_next_free_ms - now; //如果现在的时间now大于下一次可以直接获取令牌的时间，则无需等待
	// 存储的令牌有多少被使用
	double i_stored_permits_spend = MIN(permits,m_stored_permits);
	// 需要等待新生成的令牌数
	double i_fresh_permits = permits - i_stored_permits_spend;
	// 需要等待新生成的令牌数耗费时间
	__int64 i_next_free = (__int64) (i_fresh_permits*m_build_interval_ms);
	m_next_free_ms += i_next_free;
	m_stored_permits -= i_stored_permits_spend;
	
	return i_wait_ms;
}
void RateLimiter::resync(__int64 now)
{
	if (now > m_next_free_ms)
	{
		m_stored_permits =MIN((now-m_next_free_ms)/m_build_interval_ms + m_stored_permits,m_max_permits);
		m_next_free_ms = now;
	}
}

核心思想

令牌桶的思想很简单，（1）实现一个固定大小的令牌桶，按照固定的速率往桶中装入令牌，（2）每次进行流量操作（例如，下载文件）都会消耗桶中令牌，（3）当桶中令牌不足时则操作需要等待；

通过函数set_rate可以设置桶中令牌生成的平均速率；成员变量m_build_interval_ms即为每生成一个令牌需要多少时间，单位为毫秒；

从函数resync和函数claim_next中可以看出，令牌桶算法在一定程度上是能够应对突发流量，比如说，1s中生成的令牌数是1000个，而在第一秒只是用了800个令牌，那么在第2s后可使用的令牌数将为1200个。当然对于突发流量也是有上限支持的，最大为桶的大小。

acquire是限流中的核心函数，在每次进行流量操作后，需要使用acquire来决定是否需要等待，如下：

RateLimiter r;
r.set_rate(speed * 1024);
int  len = io.send(str)
if(len > 0)
	r.acquire(len);

图解流程