20220210:C++实现轻量令牌桶算法

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;
}