十年老IT知识分享 – 基于共享内存实现的令牌桶限流(带源码)

一,简述

令牌桶算法是网络流量整形和速率限制中最常使用的一种算法,关于它的描述网上也比较多资源:

 wiki:  http://en.wikipedia.org/wiki/Token_bucket

百度百科:  http://baike.baidu.com/view/2530454.htm

关于它的定义,本文不再展开过多赘述。令牌通算法标记器,视乎应用场景可以分为三种不同类型:

IN/OUT 公平标记器:流量通过令牌桶会被标记为通过(IN)和不通过(OUT)两种类型。

单速率三色标记器:有两个令牌桶(双桶出令牌速率一致),流量通过令牌桶会被标记为红黄绿三种颜色标记。业务根据实际情况对不同颜色流量作相应处理。

双速率三色标记器:跟单速率三色标记器类似,不同的地方是双桶出令牌的速率不一致。

这三种类型对应着不同的应用场景,业务根据自身特色挑选合适的标记器。

二,基于公平标记器的令牌桶算法

令牌桶算法比较简单,下面直接贴出基于公平标记器的令牌桶算法代码

Talk is cheap, show me the code!

#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; // 桶大小
};

三,在分布式系统应用令牌桶算法

从上面算法中可以看出,令牌桶算法的分布式实现关键是:保证“令牌桶”(m_fBucketSize) 和 最后变更时间(m_fLastCalcTime

)的分布式存储

而令牌桶一般要保证高性能,所以多选用类似redis这一类内存缓存。以redis为例:

1,令牌桶:保存为reids中的一个key。

2,最后变更时间:保存为reids中的一个key。

3,操作redis的时候要注意加分布式锁。

四,基于共享内存实现令牌桶算法

有一种业务场景,服务是多进程单线程模式的,这时选择基于共享内存实现令牌桶算法就比较合适了。

1,基于mmap创建共享内存。

2,基于共享内存实现一个hash table。(hash_table是为了能实现多个令牌桶,对不同类型的流量进行限流:例如针对不同ip进行限流)。

3,将令牌桶和最后变更时间保存为hash里面的一个value。

详细实现见附件源代码。

正文完