系统上为了防止瞬时流量过大造成服务和数据库崩溃，导致服务不可用，通常需要对请求限流。Go标准库中自带了限流算法的实现，即golang.org/x/time/rate，该限流器是基于令牌桶实现
令牌桶算法和漏桶算法很像：
-
漏桶限流：请求以恒定的速率放行
- 有效应对流量突刺问题
- 无法应对突发流量
-
令牌桶：每个请求需要获取到1个或多个令牌，而令牌以恒定速率放入桶中
- 如果请求速率一直大于令牌生成速率，则退化成漏桶算法
- 否则桶中的令牌会增多，下次来突发流量时能快速放行
- 能应对一定的突发流量，能应对多突发，取决于桶大小
根据描述，似乎令牌桶的实现为：
准备一个定时器和一个队列，每隔XX时间固定往队列中放入令牌
请求到来时，从队列中获取令牌，如果能获取，就执行业务逻辑，否则说明被限流了。要么等待一段时间，要么丢弃该次请求
但这么实现，需要占用一些CPU和内存资源，且桶的容量越大，占用内存资源越多
实际上常见的令牌桶算法的实现都没有用定时器和一个桶容量大小的队列，而只用4个变量：
rate：每秒生成多少令牌
capacity：桶的容量
size：当前桶中剩余令牌的数量
lastTime：上次请求令牌桶的时间
复制代码那当请求到来时，怎么根据这4个变量判断是否应该放行呢？
lastTime产生了多少令牌通过每次请求到来时再实时计算当前桶中的令牌，取代真正往桶中放令牌的过程，能有效降低对CPU和内存资源的占用
本文将介绍go rateLimiter的基本使用和底层原理分析
构造限流器可通过以下方式构造一个限流器：
limiter := rate.NewLimiter(1, 5)
复制代码- 第一个参数：每秒往桶中放多少令牌
- 第二个参数：桶的容量
rate.Every()func Every(interval time.Duration) Limit {
if interval <= 0 {
return Inf
}
return 1 / Limit(interval.Seconds())
}
复制代码go官方的实现中，初始化后桶中的令牌是满的
使请求消耗的资源和消耗的令牌数量成正比每种方法代表了当 token 不足时，不同的应对手段
reverseN入口函数为：
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation
// 简化调用ReverseN
func (lim *Limiter) Reserve() *Reservation {
return lim.ReserveN(time.Now(), 1)
}
复制代码表示在什么时间点请求（一般传time.Now()），请求多少个令牌
首先计算从当前时刻到上一次消费令牌，经过了多长时间，这期间内产生了多少令牌
tokensFromDuration// limit：每秒产生多少令牌
// d: 一段时间内
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
if limit <= 0 {
return 0
}
return d.Seconds() * float64(limit)
}
复制代码接着计算令牌桶中现在应该有多少令牌：
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
// 上传消费时间
last := lim.last
// 过去多长时间
elapsed := now.Sub(last)
// 这段时间内产生的令牌
delta := lim.limit.tokensFromDuration(elapsed)
// 加上原有的令牌数
tokens := lim.tokens + delta
// 丢弃超过容量的令牌
if burst := float64(lim.burst); tokens > burst {
tokens = burst
}
return now, last, tokens
}
复制代码计算剩余令牌减去本次消耗的令牌后，还剩余多少：
tokens -= float64(n)
复制代码如果tokens大于0，说明可以直接使用，否则就需要等待。要等待多长时间呢？为生成tokens个token，需要的时间，也是根据构造器中每秒产生多少令牌计算：
var waitDuration time.Duration
if tokens < 0 {
waitDuration = lim.limit.durationFromTokens(-tokens)
}
// 计算生成tokens个token要多长时间
func (limit Limit) durationFromTokens(tokens float64) time.Duration {
if limit <= 0 {
return InfDuration
}
seconds := tokens / float64(limit)
return time.Duration(float64(time.Second) * seconds)
}
复制代码这样如果同时有大量请求进来，这些请求就会依次排队休眠，等到了时间再结束休眠，执行业务逻辑
这样能保证先来的请求先处理
接着判断本次请求是否合法：
无论如何都无法满足ok := n <= lim.burst && waitDuration <= maxFutureReserve
复制代码如果合法，将剩余令牌数量，本次请求时间更新到限流器，再返回上面计算的信息
调用完毕，业务代码拿到返回值后可以执行如下操作：
-
OK() ：判断此次请求是否合法，是否取到了令牌
- 若不合法，一般是因为本次请求的令牌数大于桶的最大容量，需要调小请求数
-
Delay() ：若需要等待，返回等待时长
- 若大于0，说明需要在业务侧等到这么长时间再执行业务逻辑
- 否则可以立即执行
-
Cancel() ：表示取消这次的请求，将之前获取的令牌归还给限流器
第二种用法为WaitN，当桶内令牌的数量不如请求的数量时，阻塞一段时间
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)
// 简化调用
func (lim *Limiter) Wait(ctx context.Context) (err error) {
return lim.WaitN(ctx, 1)
}
复制代码可以看到参数中传入了Context，可以设置 Deadline 或者 Timeout，来决定此次 Wait 的最长时间，如果获取令牌需要等到的时间比Context的最大时间长，则该次请求不合法
也可以在外层提前取消Context，WaitN就会归还token并立即返回
select {
// 要么等待了一定时间，等到桶中有令牌
case <-t.C:
return nil
// 要么还没等到令牌，但外层执行取消操作
case <-ctx.Done():
r.Cancel()
return ctx.Err()
}
复制代码若需要配合Context中的超时设置，或cancel设置时，建议使用该限流方法
allowN第三种用法为AllowN，会立即返回是否消费成功，适用于超过限制时可以丢弃请求的场景
func (lim *Limiter) AllowN(now time.Time, n int) bool {
return lim.reserveN(now, n, 0).ok
}
func (lim *Limiter) Allow() bool {
return lim.AllowN(time.Now(), 1)
}
复制代码由于调用底层reverseN时，传的maxFutureReserve为0，则当令牌不够时就不等待，也不预先扣除令牌，对限流器无副作用
总结本文介绍了令牌桶算法的核心思想，和go官方限流器 rateLimiter的实现
漏桶算法的限流器可以参考go.uber.org/ratelimit