0%

go 令牌桶 time/rate 实现

在微服务中,为了防止在高并发场景下被大流量冲击,保护下游服务,防止服务过载,需要有控制并发请求数的手段. 限流器常用的算法有令牌桶/漏桶算法/自适应限流. golang 标准库中自带的限流算法的实现,基于令牌桶算法 rate.

令牌桶可以想象成往一个固定容量大小的桶中放 Token(可以理解成门票),系统会以固定速率往桶中放 Token,桶满了则不放.用户从桶中去 Token,如果有 Token(门票)则可以一直取. 如果没有 Token(门票)则需要等待系统发放. 令牌桶在一段时间的总 Token(门票)数是固定的,但峰值数量可以达到桶的容量.

golang time rate 的使用

初始化

1
2
3
4
5
6
7
8
9
10
// 第一个参数 r 代表每秒放入的令牌数
// 第二个参数 b 代表桶的大小
// 以下表示桶的大小为 1,每秒放入10个令牌,那么放入令牌的间隔是 0.1s
limiter := NewLimiter(10,1)

// 也可以使用另一种方法初始化,跟上面的含义一致

limit := Every(100 * time.Millisecond) // 每 100 ms生成一个token

limiter := NewLimiter(limit, 1) // 也是每秒放入10个令牌

获取 Token 方法

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
// WaitN/Wait
func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)
// WaitN 表示阻塞式获取 n 个令牌,可以传递 Context 设置超时时间
// Wait 即 WaitN(ctx,1)


// AllowN/Allow
func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool
// AllowN 表示是截止到某一时刻否能够获取到 n个令牌,满足返回 true,并消耗 n 个 Token
// 反之则返回 false,不消耗 Token
// Allow 即 AllowN(ctx,1)

// ReserveN/Reserve
func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation
// ReserveN 表示截止到某一刻获取n个token,返回一个 Reservation,提供方法判断是否需要等待多少时间
// Reserve 相当于 ReserveN(time.Now(), 1)

// ReserveN使用
// Usage example:
// r := lim.ReserveN(time.Now(), 1)
// if !r.OK() {
// // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
// return
// }
// time.Sleep(r.Delay())
// Act()

三种方法的使用场景:

  • 如果不想 drop event,可以使用 Reserve,根据返回的Reservation等待一定时间后执行
  • 如果需要遵守最后期限或取消延迟,使用 Wait
  • 如果可以允许丢弃或者忽略超过速率的事件,使用 Allow

实现原理

令牌桶原理图

看起来是有一个goroutine定时放一个队列中放入令牌,我最开始也是这样认为的,这样的好处是比较简单直观.
在 go 中是采用lazyload的方式实现的,通过记录最近一次消费时间跟当前时间进行对比更新Token数目,同时Token数目也不是使用队列,而是通过计数来实现,节省内存.

三种消费方式底层都是通过 reserveN 来实现的.
为什么能通过记录消费时间和计数来实现呢?主要是通过 Token 数可以和时间相互转化.通过limit可以得到以下信息:

  1. 生成N个Token需要多少时间
  2. 在一段时间内,可以生成多少个Token

根据以上两个信息和最近获取时间,当前获取时间来计算获取 n个 Token 需要等待的时间.

参考资料: