令牌桶限流算法的实现

📌 限流并不是为了限制业务使用，而是服务的自我保护，防止预期外的风险事件，更好的为业务服务

限流是解决各种过载、雪崩等问题的有效手段。之前了解到令牌桶限流算法的概念，直观理解有一个后台线程去生成令牌，Go官方包有令牌桶限流器的封装，令牌并非靠另一个协程异步刷入而是"读写时生成"。

限流算法

在讲令牌桶实现之前，先看看常见的限流算法有窗口算法、滑动窗口算法、漏桶算法，令牌桶算法。

计数器

计数器的思想为：在一段时间内最多允许多少请求通过。工程实践中计数器算法的实现，可以采用redis、ttl key + incrby的方式来实现窗口的限流。 计数器存在问题 计数器算法又叫固定窗口算法，在一个固定的时间，允许多少请求。计数器的实现比较简单，但有个很大的问题是没法防止边缘效应，如下图，假设阈值设置为500，请求刚好在两个窗口之间，那在4s-7s实际上允许了1000的请求量，两倍阈值的请求可能会直接打死系统。

❓ 那把窗口大小设置为1s能解决问题吗？

滑动窗口算法

为解决计数器算法的问题，引入滑动窗口的概念，限流判断的窗口不是固定的时间段，而是动态时间段的总和。比如上述的每5s是#一个窗口，在滑动窗口算法下就是"最近5s"的概念，滑动窗口算法累计最近5s的请求量，判断是否超过阈值。滑动窗口能提供相对精准的限流、一定程度解决边缘问题，缺点就是实现相对复杂，需要存储更多小窗口的状态。

漏桶算法

漏桶算法的思维：先把请求放入桶中，如果桶满则拒绝请求，桶中的请求会均匀流出。个人认为漏桶机制更适用于客户端限流，将请求匀速的发到服务端来保护下游。可类比消息队列的消费者模型，消费速率是相对恒定的。

令牌桶

令牌桶是指系统会以恒定速率生成令牌，放入令牌桶中，令牌桶满则多余的令牌丢弃，请求时需要先从令牌桶中获取令牌，获取失败则拒绝请求。

模拟令牌桶实现

按照令牌桶的定义，核心的能力

获取令牌，如果桶内有令牌，则取出一个令牌获取成功，否则获取失败
令牌定时生成，如果桶满了则丢弃多余的令牌

定义结构体TokenBucket

go 复制代码

type TokenBucket struct {
    Size int64         // 桶大小
    Unit time.Duration // 时间单位/最低1s/Unit生产Size个

    productPerSecond int64 // 每秒生成数量
    current          int64 // 当前数量

    ticker *time.Ticker // 定时器

    mu sync.Mutex
}

定时生成令牌

go 复制代码

func (tb *TokenBucket) startGenerate() {
    go func() {
        for tb.ticker != nil {
            // 定时器触发会写入该channel
            <-tb.ticker.C

            if tb.IsFull() {
                continue
            }

            tb.mu.Lock()
            if tb.current+tb.productPerSecond >= tb.Size {
                // 满了
                tb.current = tb.Size
            } else {
                tb.current = tb.current + tb.productPerSecond
            }
            tb.mu.Unlock()
        }
    }()
}

获取令牌

go 复制代码

func (tb *TokenBucket) GetToken() bool {
    tb.mu.Lock()
    defer tb.mu.Unlock()
    if tb.IsEmpty() {
        return false
    }
    tb.current = tb.current - 1
    return true
}

模拟运行的效果

time/rate实现

源码不到500行，建议大家阅读核心思想：记录上一次生成token的时间，在下次获取token的根据时间间隔计算token量。

基本使用

go 复制代码

// 参数1: 生成频率=>每s生成多少令牌
// 参数2: 桶容量
limiter := rate.NewLimiter(10, 100)

// 获取令牌成功
if limiter.Allow() {
     // 成功获取到令牌
     do()   
} else {
    return RateLimitErr
}

// 等待直到获取令牌成功 | 如果传入context超时了会返回err
err := limiter.Wait(context.Backgroud())
if err != nil {
    return err
}
// 成功获取
do()

数据结构

go 复制代码

type Limiter struct {
    mu     sync.Mutex
    limit  Limit
    burst  int
    tokens float64
    // last is the last time the limiter's tokens field was updated
    last time.Time
    // lastEvent is the latest time of a rate-limited event (past or future)
    lastEvent time.Time
}

limit: 生成频率 = 每秒生成多少token burst：桶大小 tokens: 当前token量 last: 上次生成token时间核心方法我们最主要使用的是Allow()方法，返回bool表示获取令牌结果。部分场景需要阻塞式等待会调用Wait()直至获取到token。这两方法下层的核心实现均为reserveN()，reserveN()整个方法会比较长，但是去除边界校验和结果构造，代码就很简单清晰了

go 复制代码

func (lim *Limiter) reserveN(t time.Time, n int, maxFutureReserve time.Duration) Reservation {
    // ...手动折叠
    
    t, tokens := lim.advance(t)

    // Calculate the remaining number of tokens resulting from the request.
    tokens -= float64(n)

    // Calculate the wait duration
    var waitDuration time.Duration
    if tokens < 0 {
        waitDuration = lim.limit.durationFromTokens(-tokens)
    }

    // Decide result
    ok := n <= lim.burst && waitDuration <= maxFutureReserve

    // ...手动折叠
}

advance()返回t时刻桶有多少tokens，预扣减n个，如果token不足，即扣减后的tokens < 0，则再计算需要多长时间才能满足（durationFromTokens计算生成指定token量需要多久时间）。之后判断本次获取的token的结果是否为ok，Allow()方法传递的maxFutureReserve为0，所以只要出现token不足，waitDuration必然大于0，返回的结果ok的值就是false。 advance()代码如下

go 复制代码

func (lim *Limiter) advance(t time.Time) (newT time.Time, newTokens float64) {
    // 取上一次生成的时间
    last := lim.last
    if t.Before(last) {
        last = t
    }
    // 判断时间间隔
    // Calculate the new number of tokens, due to time that passed.
    elapsed := t.Sub(last)
    // 这段时间间隔生成多少token
    delta := lim.limit.tokensFromDuration(elapsed)
    // 累计多少token
    tokens := lim.tokens + delta
    // 满了则直接用最大值
    if burst := float64(lim.burst); tokens > burst {
        tokens = burst
    }
    return t, tokens
}

获取上一次生成时间
计算时间间隔
间隔内生成多少token
累计token
判断tokens是否超过桶的大小

令牌桶的分布式限流

模拟实现的令牌桶或者Go官方包time/rate的实现，本质都是单机的限流，很少会在生产中直接使用，一般会采用分布式限流的方案。分布式限流需要依赖存储中间件来统一保存限流状态，技术选型上大部分会选择redis（高性能高可用），封装为服务或者SDK。实践一般会按照qps进行限流，因此简单的实现，可以把 "每秒生成指定量级别的qps"抽象为"每秒一个redis key"，调用incr，如果结果超过了超过了阈值qps，则表示限流，简易代码如下

go 复制代码

type TokenBucket struct {
    Key string // 标识桶
    QPS int // 阈值
}

func (tb *TokenBucket) Allow(ctx context.Context) bool {
    redisKey := func() string {
        return fmt.Sprintf("%s:%d", tb.Key, time.Now().Unix())
    }

    cmd := GetClient().Incr(ctx, redisKey())
    if cmd.Err() != nil {
        // TODO: 容灾处理
    }
    if cmd.Val() > int64(tb.QPS) {
        return false
    }
    return true
}

TokenBucket的key可以自由扩展的，比如服务限流可以考虑用【psm,cluster,method】的元组为基础生成唯一key。

和1s的窗口有什么区别？如果放在客户端，和漏桶算法有什么区别？

以上是以qps为限流对象的简易实现，标准的令牌桶实现可以参考time/rate的方案，区别点在存储介质从内存转移到数据库，因为redis大部份时候是并发安全的，所以在实现不会特别复杂。