一种更简单易懂的 token bucket 限流算法golang

2023-03-30 本文已影响0人 9_SooHyun

先给出实现的地址: https://github.com/1996Paul-Wen/watchdog

文章由

token bucket 概念简介
go的官方实现及bug说明
新实现的idea

三部分构成

1. 什么是token bucket

关于token bucket可以参考 https://mozillazg.com/2019/01/rate-limiting-intro-token-bucket.html

这里简单概括一下，token bucket是一种限流算法，主要流程是：

给出一个一定容量的令牌桶
服务提供方（server）按一定的速率往令牌桶内投放令牌，桶如果满了则无法投入新令牌
请求方（client）请求服务时需要先获取令牌

逻辑框架整体上是不复杂的，背后就是一个生产者/消费者的机制

但是，如果要考虑client等待令牌、预约令牌/取消预约等具体场景，就比较复杂了，因为里面涉及到精确的时间计算和数量计算

2. go的官方实现及bug

go对token bucket的官方实现是golang.org/x/time/rate

它完全支持client等待令牌、预约令牌/取消预约这些场景

个人看了一下源码的实现，里面关于时间点及token数量的计算逻辑比较复杂；并且，可能正是因为这里面的复杂性，导致了它不是bug free的（源码在‘取消已预约令牌’的逻辑中考虑了reserved token，是错误的），相关讨论见：

https://github.com/golang/go/issues/56924
https://stackoverflow.com/questions/70993567/rate-limiting-cancellation-token-restore
https://www.v2ex.com/t/877175
https://lailin.xyz/post/go-training-week6-3-token-bucket-2.html#comments
https://learnku.com/go/t/71323

3. 更简单易懂的算法实现

个人对token bucket算法进一步理解后，给出一种更简单易懂的算法实现

我把这种实现称为“零点（零点表示token数量为0的时间点）移动”机制，它将复杂的令牌计算直接转换成了一条时间线上简单的零点移动。下面介绍大概的思想

3.1 Token is Time

token是随着时间产生的，从这个意义上理解，token的多少就映射了时间的多少

我们说一个事件需求若干token，实际需求的是时间线上这些token对应的一段时间跨度。这个时间跨度内产生的token，都被该事件占有

从这个视角看，每个事件都以一段时间跨度的形式分布在时间线上，并且跨度之间互不重合

如果把token is time这个概念可视化出来，就像：

    ...  >|< event1 occupied this time span, which generates 3 tokens   >|<   event2 occupied this span, with 2 t >|<  ...  >
----------|--------------------|--------------------|--------------------|--------------------|--------------------|---> timeline

3.2 OCCUPIED TIME SPAN

OCCUPIED TIME SPAN 表示事件对时间线上一个时间跨度的占有。一个事件持有一个OCCUPIED TIME SPAN，表示该时间跨度内产生的所有token都将归该事件使用

3.3 DEPRECATED TIME SPAN

因为bucket的容量有限，如果一段较长的时间内没有新事件，那么这段时间内生产出的所有token中，超出bucket容量的那部分token会被丢弃（deprecated tokens）

deprecated tokens 对应的时间跨度称为 DEPRECATED TIME SPAN，这个跨度内的时间都认为是无用的

注：我们可以将bucket视为队列，那些DEPRECATED TIME SPAN内产生的tokens根据先进先出原则都被出队丢弃了

3.4 ZERO POINT 零点

现在，让我们关注标志着当前时间线上最近一个【新的、可用的时间跨度】开始的时间点。我们称这个时间点为ZERO POINT。这是因为此刻可申请的tokens总是为0，而新事件的 OCCUPIED TIME SPAN 最早只能从该点开始

例如：

                    ZERO POINT  
...-----------------------|---------------------------------|----------------|->  timeline
...last OccupiedTimeSpan >|<  new time span to  occupy     >|

3.5 ZERO POINT movement logic

那么，零点该如何移动呢？也很简单，因为只有前后移两种情况：

前移（向时间线正方向移动）
- 新事件占用token。新事件占用token就是占用相应的时间跨度，也就是将零点向前移动相同长度的OCCUPIED TIME SPAN
- 出现了DEPRECATED TIME SPAN。很明显，DEPRECATED TIME SPAN的出现会将零点向前移动相同长度的DEPRECATED TIME SPAN
后移（向时间线负方向移动）
- 事件在发生前释放它所占用的token。释放占用的token就是取消占用的时间跨度，也就是将零点向后移动相同长度的占用时间跨度。这是将零点向后移动的唯一方法

基于上面的想法，将零点移动机制做了个实现，叫watchdog：https://github.com/1996Paul-Wen/watchdog

4. time/rate 和 watchdog 对比

watchdog更简单

将time/rate和watchdog的Limiter结构体对比，可见一斑

// ---------- watchdog's Limiter struct------------
type Limiter struct {
    limit float64 
    burst float64 

    zeroPoint time.Time  // zeroPoint marks the start of a new useful time span
    mu        sync.Mutex // protects zeroPoint from concurrency access
}


// ---------- time/rate's Limiter struct------------
type Limiter struct {
    limit  Limit
    burst  int

    mu     sync.Mutex
    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
}

time/rate的 Limiter 维护了tokens字段，表示在last时刻 tokens 的数量。同时还维护了一个lastEvent，表示一个最远事件的发生时间。那么在这几个数据之上，就免不了时间-数量对应关系的复杂计算

而watchdog的 Limiter 只维护了一个零点，在等待令牌、预约令牌/取消预约等各种可能场景下，都只需要修改一个值，那就是zeroPoint。建立在一个数据上的计算逻辑，会更加简单，因而更加可靠、可维护

另外，watchdog的用法可以covertime/rate，它保持了相似的函数签名和语义，因此使用起来无需过多的心智负担

更多信息可以直接到https://github.com/1996Paul-Wen/watchdog