令牌桶限速C++实现

由于业务需求，要对总流量进行限速，做了一个C++版本的实现，基本满足业务要求。有需要的同学自行进行优化。
a. 按特定的速率向令牌桶投放令牌
b. 当桶中有足够的令牌则报文可以被继续发送下去，同时令牌桶中的令牌量按报文的长度做相应的减少；
c. 当令牌桶中的令牌不足时，报文将不能被发送，只有等到桶中生成了新的令牌，报文才可以发送。这就可以限制报文的流量只能是小于等于令牌生成的速度，达到限制流量的目的。

生产者按特定的速率向令牌桶投放令牌；
消费者查看令牌桶的令牌，令牌充足消费相应数量的令牌，令牌不足丢弃相应报文或请求。

瞬时流量过高可以用限流处理，流量一直过载的话要考虑扩容。

class ClimitSpeed
{
public:
    ClimitSpeed(unsigned int speed)
    {
        m_bExitFlag  = false;
        m_bStartFlag = false;
        m_time  = 1000;
        m_totalByte = speed;
        m_sumByte = 0;
        InitializeCriticalSection(&m_rwlock);
    }
    ~ClimitSpeed()
    {
        m_bExitFlag = true;
        DeleteCriticalSection(&m_rwlock);
    }

    void Start()
    {
        if (m_bStartFlag)
        {
            return;
        }

        m_bStartFlag = true;
        
        _beginthread(threadFeed, 0, this);
    }
    bool CanConsume()
    {
        return (m_sumByte > 0);
    }
    void Consume(unsigned int len)
    {
        EnterCriticalSection(&m_rwlock);   
        if (m_sumByte > len)
        {
            m_sumByte -= len;
        } 
        else
        {
            m_sumByte = 0;
        }
        LeaveCriticalSection(&m_rwlock);
    }
private:
    static void threadFeed(void* arg);
    void Feed()
    {
        EnterCriticalSection(&m_rwlock);    
        m_sumByte += m_totalByte;
        LeaveCriticalSection(&m_rwlock);

    }
    bool IsStart()
    {
        return m_bStartFlag;
    }
private:
    unsigned int m_time;  
    unsigned int m_sumByte;
    unsigned int m_totalByte;
    bool m_bStartFlag;
    bool m_bExitFlag;
    CRITICAL_SECTION m_rwlock;
};

测试代码

//////////////////////////////////////////////////////////////////////////
/*
在入口函数中包含 _CrtDumpMemoryLeaks();  
即可检测到内存泄露
*/
#define CRTDBG_MAP_ALLOC  
#include <stdlib.h>  
#include <crtdbg.h>  
//#ifdef _DEBUG  
//#define new new(_NORMAL_BLOCK, __FILE__, __LINE__)  
//#endif  
//////////////////////////////////////////////////////////////////////////
#include "stdafx.h"
#include <math.h>
#include <conio.h>    //for kbhit()
#include <stdint.h>
#include <time.h>
#include <process.h>
#include <Windows.h>
static void EnableMemLeakCheck()  
{  
    int tmpFlag = _CrtSetDbgFlag(_CRTDBG_REPORT_FLAG);  
    tmpFlag |= _CRTDBG_LEAK_CHECK_DF;  
    _CrtSetDbgFlag(tmpFlag);  
}  

class ClimitSpeed
{
public:
    ClimitSpeed(unsigned int speed)
    {
        m_bExitFlag  = false;
        m_bStartFlag = false;
        m_time  = 1000;
        m_totalByte = speed;
        m_sumByte = 0;
        InitializeCriticalSection(&m_rwlock);
    }
    ~ClimitSpeed()
    {
        m_bExitFlag = true;
        DeleteCriticalSection(&m_rwlock);
    }

    void Start()
    {
        if (m_bStartFlag)
        {
            return;
        }

        m_bStartFlag = true;
        
        _beginthread(threadFeed, 0, this);
    }
    bool CanConsume()
    {
        return (m_sumByte > 0);
    }
    void Consume(unsigned int len)
    {
        EnterCriticalSection(&m_rwlock);   
        if (m_sumByte > len)
        {
            m_sumByte -= len;
        } 
        else
        {
            m_sumByte = 0;
        }
        LeaveCriticalSection(&m_rwlock);
    }
private:
    static void threadFeed(void* arg);
    void Feed()
    {
        EnterCriticalSection(&m_rwlock);    
        m_sumByte += m_totalByte;
        LeaveCriticalSection(&m_rwlock);

    }
    bool IsStart()
    {
        return m_bStartFlag;
    }
private:
    unsigned int m_time;  
    unsigned int m_sumByte;
    unsigned int m_totalByte;
    bool m_bStartFlag;
    bool m_bExitFlag;
    CRITICAL_SECTION m_rwlock;
};

void ClimitSpeed::threadFeed(void* arg)
{
    ClimitSpeed* pLimit = (ClimitSpeed*)arg;
    while(!pLimit->m_bExitFlag)
    {
        if (pLimit->m_bStartFlag)
        {
            pLimit->Feed();
        }
        Sleep(pLimit->m_time);
    }
}

//生成随机数，模拟接受随机报文长度
int GetIntLH(int low, int hight)
{
    static int a = 0x100;
    int seed = rand();
    srand(static_cast<unsigned>(time(nullptr)) - seed + (a += 50));
    return seed % (hight - low + 1) + low;
}
typedef struct
{
    bool* pExitFlag;
    unsigned int speed;
    HANDLE thread;
}StrucThreadConsume;

//消费者线程
static UINT __stdcall funcrun(void* arg)
{
    int msgLen = 0;
    unsigned int sum = 0;

    StrucThreadConsume* p = (StrucThreadConsume *)arg;

    //测试生成限流器 B/s
    ClimitSpeed *pclimit = new ClimitSpeed(p->speed);

    unsigned int tick = GetTickCount();
    unsigned int tick2 = GetTickCount();

    pclimit->Start();

    while(!(*(p->pExitFlag)))
    {
        tick2 = GetTickCount();
        if (tick2 - tick >= 1000)
        {
            printf("thread:%#x, sum:%lu, time:%lu\n", p->thread, sum, tick2 - tick);
            sum = 0;
            tick = tick2;
        }

        if (!pclimit->CanConsume())
        {
            Sleep(2);
            continue;
        }

        msgLen = GetIntLH(64, 1500);

        pclimit->Consume(msgLen);

        sum += msgLen;
    }
    delete pclimit;
    pclimit = NULL;
    printf("thread:%#x funcrun exit\n", p->thread);
    return 0;
}



int main()
{
    EnableMemLeakCheck();
    //_CrtSetBreakAlloc(265); 

    bool bExit = false;

    StrucThreadConsume t1;
    t1.pExitFlag = &bExit;
    t1.speed = 20 * 1024 * 1024;

    StrucThreadConsume t2;
    t2.pExitFlag = &bExit;
    t2.speed = 30 * 1024 * 1024;

    HANDLE thread1 = (HANDLE)_beginthreadex(NULL, 0, funcrun, &t1, CREATE_SUSPENDED, NULL);
    HANDLE thread2 = (HANDLE) _beginthreadex(NULL, 0, funcrun, &t2, CREATE_SUSPENDED, NULL);

    t1.thread = thread1;
    t2.thread = thread2;

    ResumeThread(thread1);
    ResumeThread(thread2);

    char k = 0;
    while(1)
    {
        if (_kbhit())
        {
            k = _getch();
            if(k == 'q')
            {
                bExit = true;
                break;
            }
        }

    }

    printf("main will exit\n");

    getchar();
    return 0;
}

测试截图，两个测试线程流量控制在20MB/s和30MB/s