基于共享内存实现的令牌桶限流(带源码)

#pragma once
#include <sys/time.h>
#include <stdint.h>
#pragma pack(push, 1)

class TokenBucket
{
public:
    TokenBucket() :m_fLastCalcTime(0), m_fBucketWater(0), m_fRateUnitsPerSeconds(0), m_fBucketSize(0)
    {}

    TokenBucket(double fRateUnitsPerSeconds, double fBucketSize) :m_fLastCalcTime(0), m_fBucketWater(0)
    {
        Set(fRateUnitsPerSeconds, fBucketSize);
    }
    
    // 设置速率和桶大小
    void Set(double fRateUnitsPerSeconds, double fBucketSize)
    {
        m_fRateUnitsPerSeconds = fRateUnitsPerSeconds;
        m_fBucketSize = fBucketSize;
    }

    // 按照当前时间更新桶水量
    void UpdateBucketWater()
    {
        struct timeval tvNow;
        gettimeofday(&tvNow, NULL);
        double fNow = (double)tvNow.tv_sec + (double)tvNow.tv_usec / 1000000.0;
        double fWaterFlowed = m_fRateUnitsPerSeconds * ( fNow - m_fLastCalcTime );
        m_fBucketWater += fWaterFlowed;
        if ( m_fBucketWater > m_fBucketSize ) m_fBucketWater = m_fBucketSize;
        m_fLastCalcTime = fNow;
    }

    // 判断能否发送指定的量，能就从桶中减去这些量，否则调用者就抛弃或Sleep等待到能发的时候再发
    bool CanSend(double fSendUnits)
    {
        UpdateBucketWater();
        if (m_fBucketWater >= fSendUnits) {
            m_fBucketWater -= fSendUnits;
            return true;
        } else {
            return false;
        } 
    }

    // 当调用CanSend()判断可以发送后，实际上又没有发送那么多，就需要调用Compensate()把未发送出去的量补回到桶中，这样流量控制才会准确
    void Compensate(double fSendLeftUnits)
    {
        m_fBucketWater += fSendLeftUnits;
    }

    //下面是后判断模式专用：用于判断当时时间桶中水量是否为负数,不为负数表示可以接着干活
    bool CanGo()
    {
        UpdateBucketWater();
        if (m_fBucketWater >= 0) {
            return true;
        }
        else {
            return false;
        } 
    }
    //活干完才知道干了多少，此处就直接从桶中扣掉，不管是正是负,后付费模式
    void AfterPayment(double fCost)
    {
        m_fBucketWater -= fCost;
    }
    
public:
    double m_fLastCalcTime; // 最后更新时间
    double m_fBucketWater; // 桶中水量，不能超过m_fBucketSize
    double m_fRateUnitsPerSeconds; // 速率：每秒放多少个
    double m_fBucketSize; // 桶大小
};