Golang限流器time/rate正确打开姿势

今天聊一下go语言限流工具的 golang.org/x/time/rate 包下Limiter的用法

用Limiter做一个qps限流器

我用这个限流工具做了一个qps限流的功能。

假设我限制qps为5，创建一个Limiter。

go 代码解读复制代码    limiter := rate.NewLimiter(rate.Limit(5), 5)

注意按Limiter的文档描述，它是一个令牌桶（token bucket）。

上面这个NewLimiter第一个参数是limit，表示每秒生成token的速率。第二个参数是burst，表示一次最多可以获得多少个token。

然后调用Allow()，尝试消费掉一个token，结果true为成功，false则表示无法获得token，应该被限流。

go 代码解读复制代码if limit.Allow() {
    ...
}

至此一切正常。

但，当我某天想动态调整qps，于是我调用SetLimiter(10)，意图把后续的qps改为10。

go 代码解读复制代码    limiter.SetLimit(10)

但发现，接下来的第二秒，Allow只能得到5次为true的结果？！

Limiter的工作机制

网上没有明确的解释。

遇到这种没有什么捷径，只能自己读一遍代码了。

不读不知道，这个库实现很有意思，它有一个小而美（简单但好用）算法。

但也有令人感到迷惑的外表。

结构定义

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
}

这个结构的字段越看越奇怪，tokens（令牌）为什么是一个float64类型的值？难道可以有0.1个令牌？

仔细看，还真是！！

算法：令牌桶 or 滑动窗口 ？

网上都说它是一个令牌桶的算法，它的文档也是这样描述的。

可是仔细看，我认为它只有外表（接口）像是令牌桶。

实际上它是一个披着令牌桶外皮的滑动窗口算法。

它的算法可以用一句话描述为：

每次调用时，按时间计算可用的token数量。

image.png

核心是它的advance函数，它负责根据limit状态和时间重新计算limit的状态。

go 代码解读复制代码func (lim *Limiter) advance(t time.Time) (newT time.Time, newTokens float64) {
    last := lim.last
    if t.Before(last) {
       last = t
    }
    
    elapsed := t.Sub(last) // 从上次调用过去了多久
    delta := lim.limit.tokensFromDuration(elapsed)  // 过去的时间包含多少token
    tokens := lim.tokens + delta  // 可用的token数
    if burst := float64(lim.burst); tokens > burst {
       tokens = burst
    }
    return t, tokens
}

func (limit Limit) durationFromTokens(tokens float64) time.Duration {
    if limit <= 0 {
       return InfDuration
    }
    seconds := tokens / float64(limit)
    return time.Duration(float64(time.Second) * seconds)
}

下面这个时间流图，简单展示了一个限流器的工作过程，因为qps为5，一个方格表示过去200ms。

image.png

1. Limiter第一次调用Allow()，算出token数为5（初始化brust），消费1个token，返回true。
2. Limiter第二次调用AllowN(4)，因为离上次调用时间1.4秒，可用token数为5（超过burst的被忽略），消费4个token，返回true。
3. Limiter第三次调用AllowN(5)，离上次调用550ms左右，算出可用token数为3.75个，将要消费5个，token不足，返回false

只能从代码里意会的设计

Limiter假定token是以恒定速率生成的，如qps为5，所以每200ms可以得到一个token。

每次调用的时间不固定，token也可以获得小数个，比如过去20ms调用，会发现token数增加了0.1。

这种方法，我把它总结为

把token空间遇映射到时间轴上

优点

空间复杂度极低，O(1)，不管你希望每秒产生多少个token，Limiter始终只有5个变量.

时间复杂度极低，O(1)，每次调用时一步算出token数量，没有额外的开销。

由于把token映射到时间轴上，可以很方便地知道需要的token什么时候能生成，这就让Wait()和WaitN()的实现变得非常简单。

另外，由于token数量是一个非常直观的变量，可以随时打印出来看看，非常方便定位问题。

小巧却强大的一个工具。

缺点

从文档上不容易看出来Limit和Burst是怎么被使用的。

我甚至认为它们俩可以合成一个初始化参数的。

最后：修正开头的问题

文章一开头的qps限流工作不正常的问题，应该做如下修改

go 代码解读复制代码    limiter.SetLimit(10)
    limiter.SetBurst(10)

既增加每秒生成token的速度（limit）又增加桶的大小（burst）