C++超輕量限流器實現

目錄

一、簡介

二、令牌桶算法簡述

三、使用

四、實現

五、後續


一、簡介

       該限流器基於令牌桶算法實現,特點如下:

  1. 接入方便,無業務侵入:接入只需要添加一行代碼r.pass()。
  2. 線程安全,CPU友好,無鎖高效。
  3. 輕量,核心代碼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、提供定長的有序等待隊列。

       最後,有問題都可以私信我,畢竟人無完人,本文有地方描述得不太準確希望能夠告知。

 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章