目錄
一、簡介
該限流器基於令牌桶算法實現,特點如下:
- 接入方便,無業務侵入:接入只需要添加一行代碼r.pass()。
- 線程安全,CPU友好,無鎖高效。
- 輕量,核心代碼200行。
github地址,有興趣可以來顆star:https://github.com/YukangLiu/RateLimiter
二、令牌桶算法簡述
有一個裝着令牌(token)的桶,它按照qps的速率補充令牌,滿了則溢出。當一個請求到來時,它需要先從桶中取出一個令牌,才能繼續進行後續操作。所以簡單來說,令牌桶算法是基於令牌的補給速率限制qps。
三、使用
先來看看如何使用,其使用非常簡單,先創建一個限流器對象,然後在需要限流的地方調用pass()方法即可,示例如下:
int main()
{
RateLimiter r(100);//100qps限流器
for(int i = 0; i < 500; ++i)
{
r.pass();//通過的速率爲100/s
}
return 0;
}
這個for循環執行完需要5s。
四、實現
//限流器
class RateLimiter
{
public:
//qps限制最大爲十億
RateLimiter(int64_t qps);
DISALLOW_COPY_MOVE_AND_ASSIGN(RateLimiter);
//對外接口,能返回說明流量在限定值內
void pass();
private:
//獲得當前時間,單位ns
int64_t now();
//更新令牌桶中的令牌
void supplyTokens();
//嘗試獲得令牌
bool tryGetToken();
//必定成功獲得令牌
void mustGetToken();
//令牌桶大小
const int64_t bucketSize_;
//剩下的token數
AtomicSequence tokenLeft_;
//補充令牌的單位時間
const int64_t supplyUnitTime_;
//上次補充令牌的時間,單位納秒
int64_t lastAddTokenTime_;
//自旋鎖
Spinlock lock_;
};
該類主要維護的是令牌桶的容量、當前桶中還存在的令牌數、上次補充令牌的時間這三個變量。其中AtomicSequence就是對原子變量的一個封裝,主要爲了防止false sharing,其具體設計可以看我另一篇博客Disruptor原理概述與輕量級C++實現
具體實現如下:
int64_t RateLimiter::now()
{
struct timeval tv;
::gettimeofday(&tv, 0);
int64_t seconds = tv.tv_sec;
return seconds * NS_PER_SECOND + tv.tv_usec * NS_PER_USECOND;
}
上述方法用於獲取當前時間的納秒數。
void RateLimiter::supplyTokens()
{
auto cur = now();
if (cur - lastAddTokenTime_ < supplyUnitTime_)
{
return;
}
{
SpinlockGuard lock(lock_);
//等待自旋鎖期間可能已經補充過令牌了
int64_t newTokens = (cur - lastAddTokenTime_) / supplyUnitTime_;
if (newTokens <= 0)
{
return;
}
//更新補充時間,不能直接=cur,否則會導致時間丟失
lastAddTokenTime_ += (newTokens * supplyUnitTime_);
auto freeRoom = bucketSize_ - tokenLeft_.load();
if(newTokens > freeRoom || newTokens > bucketSize_)
{
newTokens = freeRoom > bucketSize_ ? bucketSize_ : freeRoom;
}
tokenLeft_.fetch_add(newTokens);
}
}
上述方法用於補充令牌:計算上次補充到當前過了多久,然後根據qps補充相應數量的令牌,最後更新補充令牌的時間。這裏有一個細節就是更新的時間並不是當前的時間,而是補充相應令牌數對應需要經過的時間。比如1s要補充1個令牌,本次補充距上次過了1.8s,此時只補充1個令牌,補充令牌的時間只比上次增加1s而不是增加1.8s。
另外,這裏用到了自旋鎖,該自旋鎖是用原子變量模擬信號量來實現的,不會讓進程睡眠,這裏不多介紹,有興趣可以到github查看源碼。
bool RateLimiter::tryGetToken()
{
supplyTokens();
//獲得一個令牌
auto token = tokenLeft_.fetch_add(-1);
if(token <= 0)
{//已經沒有令牌了,歸還透支的令牌
tokenLeft_.fetch_add(1);
return false;
}
return true;
}
上述方法用於嘗試獲取令牌,首先回去補充令牌桶,然後取一個令牌,如果桶中已沒有令牌則返回失敗。
void RateLimiter::mustGetToken()
{
bool isGetToken = false;
for(int i = 0; i < RETRY_IMMEDIATELY_TIMES; ++i)
{
isGetToken = tryGetToken();
if(isGetToken)
{
return;
}
}
while(1)
{
isGetToken = tryGetToken();
if(isGetToken)
{
return;
}
else
{
//讓出CPU
sleep(0);
}
}
}
上述方法用於必須獲取令牌,不獲得就不返回。它在tryGetToken的基礎上增加了retry機制:首先自旋retry一定次數,還是沒拿到令牌說明要麼此時併發量很大,要麼qps較小,需要較長時間才能補充令牌,所以這之後每次retry失敗都會讓出cpu,讓cpu先去處理其它任務而不是佔用CPU資源。所以這種retry機制是對CPU友好的,但是同時也不會導致併發量下降。
void RateLimiter::pass()
{
return mustGetToken();
}
該方法則是對外接口,目前是直接調用mustGetToken方法。
五、後續
這個工具接下來可以延申的地方有1、開放tryPass接口。2、提供定長的有序等待隊列。
最後,有問題都可以私信我,畢竟人無完人,本文有地方描述得不太準確希望能夠告知。