go-zero 是如何实现令牌桶限流的？

原文链接：
上一篇文章介绍了 如何实现计数器限流？主要有两种实现方式，分别是固定窗口和滑动窗口，并且分析了 go-zero 采用固定窗口方式实现的源码。
但是采用固定窗口实现的限流器会有两个问题：

• 会出现请求量超出限制值两倍的情况
  
• 无法很好处理流量突增问题
  

这篇文章来介绍一下令牌桶算法，可以很好解决以上两个问题。
工作原理

算法概念如下：

• 令牌以固定速率生成；
  
• 生成的令牌放入令牌桶中存放，如果令牌桶满了则多余的令牌会直接丢弃，当请求到达时，会尝试从令牌桶中取令牌，取到了令牌的请求可以执行；
  
• 如果桶空了，那么尝试取令牌的请求会被直接丢弃。
  

令牌桶算法既能够将所有的请求平均分布到时间区间内，又能接受服务器能够承受范围内的突发请求，因此是目前使用较为广泛的一种限流算法。
源码实现

源码分析我们还是以 go-zero 项目为例，首先来看生成令牌的部分，依然是使用 Redis 来实现。

1. // core/limit/tokenlimit.go
3. // 生成 token 速率
4. script = `local rate = tonumber(ARGV[1])
5. // 通容量
6. local capacity = tonumber(ARGV[2])
7. // 当前时间戳
8. local now = tonumber(ARGV[3])
9. // 请求数量
10. local requested = tonumber(ARGV[4])
11. // 需要多少秒才能把桶填满
12. local fill_time = capacity/rate
13. // 向下取整，ttl 为填满时间 2 倍
14. local ttl = math.floor(fill_time*2)
15. // 当前桶剩余容量，如果为 nil，说明第一次使用，赋值为桶最大容量
16. local last_tokens = tonumber(redis.call("get", KEYS[1]))
17. if last_tokens == nil then
18.     last_tokens = capacity
19. end
21. // 上次请求时间戳，如果为 nil 则赋值 0
22. local last_refreshed = tonumber(redis.call("get", KEYS[2]))
23. if last_refreshed == nil then
24.     last_refreshed = 0
25. end
27. // 距离上一次请求的时间跨度
28. local delta = math.max(0, now-last_refreshed)
29. // 距离上一次请求的时间跨度能生成的 token 数量和桶内剩余 token 数量的和
30. // 与桶容量比较，取二者的小值
31. local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
32. // 判断请求数量和桶内 token 数量的大小
33. local allowed = filled_tokens >= requested
34. // 被请求消耗掉之后，更新剩余 token 数量
35. local new_tokens = filled_tokens
36. if allowed then
37.     new_tokens = filled_tokens - requested
38. end
40. // 更新 redis token
41. redis.call("setex", KEYS[1], ttl, new_tokens)
42. // 更新 redis 刷新时间
43. redis.call("setex", KEYS[2], ttl, now)
45. return allowed`

复制代码

Redis 中主要保存两个 key，分别是 token 数量和刷新时间。
核心思想就是比较两次请求时间间隔内生成的 token 数量 + 桶内剩余 token 数量，和请求量之间的大小，如果满足则允许，否则则不允许。
限流器初始化：

1. // A TokenLimiter controls how frequently events are allowed to happen with in one second.
2. type TokenLimiter struct {
3.     // 生成 token 速率
4.     rate           int
5.     // 桶容量
6.     burst          int
7.     store          *redis.Redis
8.     // 桶 key
9.     tokenKey       string
10.     // 桶刷新时间 key
11.     timestampKey   string
12.     rescueLock     sync.Mutex
13.     // redis 健康标识
14.     redisAlive     uint32
15.     // redis 健康监控启动状态
16.     monitorStarted bool
17.     // 内置单机限流器
18.     rescueLimiter  *xrate.Limiter
19. }
21. // NewTokenLimiter returns a new TokenLimiter that allows events up to rate and permits
22. // bursts of at most burst tokens.
23. func NewTokenLimiter(rate, burst int, store *redis.Redis, key string) *TokenLimiter {
24.     tokenKey := fmt.Sprintf(tokenFormat, key)
25.     timestampKey := fmt.Sprintf(timestampFormat, key)
27.     return &TokenLimiter{
28.         rate:          rate,
29.         burst:         burst,
30.         store:         store,
31.         tokenKey:      tokenKey,
32.         timestampKey:  timestampKey,
33.         redisAlive:    1,
34.         rescueLimiter: xrate.NewLimiter(xrate.Every(time.Second/time.Duration(rate)), burst),
35.     }
36. }

复制代码

其中有一个变量 rescueLimiter，这是一个进程内的限流器。如果 Redis 发生故障了，那么就使用这个，算是一个保障，尽量避免系统被突发流量拖垮。

提供了四个可调用方法：

1. // Allow is shorthand for AllowN(time.Now(), 1).
2. func (lim *TokenLimiter) Allow() bool {
3.     return lim.AllowN(time.Now(), 1)
4. }
6. // AllowCtx is shorthand for AllowNCtx(ctx,time.Now(), 1) with incoming context.
7. func (lim *TokenLimiter) AllowCtx(ctx context.Context) bool {
8.     return lim.AllowNCtx(ctx, time.Now(), 1)
9. }
11. // AllowN reports whether n events may happen at time now.
12. // Use this method if you intend to drop / skip events that exceed the rate.
13. // Otherwise, use Reserve or Wait.
14. func (lim *TokenLimiter) AllowN(now time.Time, n int) bool {
15.     return lim.reserveN(context.Background(), now, n)
16. }
18. // AllowNCtx reports whether n events may happen at time now with incoming context.
19. // Use this method if you intend to drop / skip events that exceed the rate.
20. // Otherwise, use Reserve or Wait.
21. func (lim *TokenLimiter) AllowNCtx(ctx context.Context, now time.Time, n int) bool {
22.     return lim.reserveN(ctx, now, n)
23. }

复制代码

最终调用的都是 reverveN 方法：

1. func (lim *TokenLimiter) reserveN(ctx context.Context, now time.Time, n int) bool {
2.     // 判断 Redis 健康状态，如果 Redis 故障，则使用进程内限流器
3.     if atomic.LoadUint32(&lim.redisAlive) == 0 {
4.         return lim.rescueLimiter.AllowN(now, n)
5.     }
7.     // 执行限流脚本
8.     resp, err := lim.store.EvalCtx(ctx,
9.         script,
10.         []string{
11.             lim.tokenKey,
12.             lim.timestampKey,
13.         },
14.         []string{
15.             strconv.Itoa(lim.rate),
16.             strconv.Itoa(lim.burst),
17.             strconv.FormatInt(now.Unix(), 10),
18.             strconv.Itoa(n),
19.         })
20.     // redis allowed == false
21.     // Lua boolean false -> r Nil bulk reply
22.     if err == redis.Nil {
23.         return false
24.     }
25.     if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
26.         logx.Errorf("fail to use rate limiter: %s", err)
27.         return false
28.     }
29.     if err != nil {
30.         logx.Errorf("fail to use rate limiter: %s, use in-process limiter for rescue", err)
31.         // 如果有异常的话，会启动进程内限流
32.         lim.startMonitor()
33.         return lim.rescueLimiter.AllowN(now, n)
34.     }
36.     code, ok := resp.(int64)
37.     if !ok {
38.         logx.Errorf("fail to eval redis script: %v, use in-process limiter for rescue", resp)
39.         lim.startMonitor()
40.         return lim.rescueLimiter.AllowN(now, n)
41.     }
43.     // redis allowed == true
44.     // Lua boolean true -> r integer reply with value of 1
45.     return code == 1
46. }

复制代码

最后看一下进程内限流的启动与恢复：

1. func (lim *TokenLimiter) startMonitor() {
2.     lim.rescueLock.Lock()
3.     defer lim.rescueLock.Unlock()
5.     // 需要加锁保护，如果程序已经启动了，直接返回，不要重复启动
6.     if lim.monitorStarted {
7.         return
8.     }
10.     lim.monitorStarted = true
11.     atomic.StoreUint32(&lim.redisAlive, 0)
13.     go lim.waitForRedis()
14. }
16. func (lim *TokenLimiter) waitForRedis() {
17.     ticker := time.NewTicker(pingInterval)
18.     // 更新监控进程的状态
19.     defer func() {
20.         ticker.Stop()
21.         lim.rescueLock.Lock()
22.         lim.monitorStarted = false
23.         lim.rescueLock.Unlock()
24.     }()
26.     for range ticker.C {
27.         // 对 redis 进行健康监测，如果 redis 服务恢复了
28.         // 则更新 redisAlive 标识，并退出 goroutine
29.         if lim.store.Ping() {
30.             atomic.StoreUint32(&lim.redisAlive, 1)
31.             return
32.         }
33.     }
34. }

复制代码

以上就是本文的全部内容，如果觉得还不错的话欢迎点赞，转发和关注，感谢支持。
参考文章：

• https://juejin.cn/post/7052171117116522504
  
• https://www.infoq.cn/article/Qg2tX8fyw5Vt-f3HH673
  

推荐阅读：

• 如何实现计数器限流？
  
• go-zero 是如何做路由管理的？
  

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！