1.RateLimiter.h
#ifndef _RATE_LIMITER_H_
#define _RATE_LIMITER_H_
#include <mutex>
#define RETRY_IMEDIATELY_TIMES 30
class RateLimiter {
public:
RateLimiter(int qps);
void pass();
private:
//互斥锁
std::mutex mtx;
//获得当前时间, 单位ns
long now();
//更新令牌桶中的令牌
void supplyTokens();
//尝试获得令牌
bool tryGetToken();
//必定成功获得令牌
void mustGetToken();
//令牌桶大小
int bucketSize;
//剩下的令牌数
int tokenLeft;
//补充令牌的单位时间
long supplyUnitTime;
//上次补充令牌的时间
long lastAddTokenTime;
};
#endif
2.RateLimiter.cpp
#include <iostream>
#include <sys/types.h>
#include <sys/time.h>
#include <unistd.h>
#include <time.h>
#include "RateLimiter.h"
using namespace std;
RateLimiter::RateLimiter(int qps)
{
supplyUnitTime = 1000000000L / qps;
bucketSize = 100;
tokenLeft = 0;
}
long RateLimiter::now()
{
struct timeval tv;
long seconds;
gettimeofday(&tv, 0);
seconds = tv.tv_sec;
return seconds * 1000000000L + tv.tv_usec * 1000L;
}
void RateLimiter::supplyTokens()
{
long cur = now();
if(cur - lastAddTokenTime < supplyUnitTime)
{
return;
}
long newTokens = (cur - lastAddTokenTime) / supplyUnitTime;
if(newTokens <= 0)
{
return;
}
mtx.lock();
//更新补充时间
lastAddTokenTime += (newTokens * supplyUnitTime);
int freeRoom = bucketSize - tokenLeft;
if(newTokens > freeRoom || freeRoom > bucketSize)
{
newTokens = freeRoom > bucketSize ? bucketSize : freeRoom;
}
tokenLeft += newTokens;
mtx.unlock();
}
bool RateLimiter::tryGetToken()
{
int token;
supplyTokens();
mtx.lock();
token = tokenLeft;
tokenLeft--;
mtx.unlock();
if(token <= 0)
{
mtx.lock();
tokenLeft++;
mtx.unlock();
return false;
}
return true;
}
void RateLimiter::mustGetToken()
{
bool isGetToken = false;
for(int i=0; i<RETRY_IMEDIATELY_TIMES; i++)
{
isGetToken = tryGetToken();
if(isGetToken)
{
return;
}
}
while(1)
{
isGetToken = tryGetToken();
if(isGetToken)
{
return;
}
else
{
sleep(0);
}
}
}
void RateLimiter::pass()
{
return mustGetToken();
}
int main()
{
RateLimiter r(1000);
for(int i=0; i<500; i++)
{
r.pass();
//printf("pass %d\n", i);
}
return 0;
}