go-zero 是如何实现计数器限流的？

原文链接： 如何实现计数器限流？
上一篇文章 go-zero 是如何做路由管理的？ 介绍了路由管理，这篇文章来说说限流，主要介绍计数器限流算法，具体的代码实现，我们还是来分析微服务框架 go-zero 的源码。
在微服务架构中，一个服务可能需要频繁地与其他服务交互，而过多的请求可能导致性能下降或系统崩溃。为了确保系统的稳定性和高可用性，限流算法应运而生。
限流算法允许在给定时间段内，对服务的请求流量进行控制和调整，以防止资源耗尽和服务过载。
计数器限流算法主要有两种实现方式，分别是：

• 固定窗口计数器
  
• 滑动窗口计数器
  

下面分别来介绍。
固定窗口计数器

算法概念如下：

• 将时间划分为多个窗口；
  
• 在每个窗口内每有一次请求就将计数器加一；
  
• 如果计数器超过了限制数量，则本窗口内所有的请求都被丢弃当时间到达下一个窗口时，计数器重置。
  

固定窗口计数器是最为简单的算法，但这个算法有时会让通过请求量允许为限制的两倍。

考虑如下情况：限制 1 秒内最多通过 5 个请求，在第一个窗口的最后半秒内通过了 5 个请求，第二个窗口的前半秒内又通过了 5 个请求。这样看来就是在 1 秒内通过了 10 个请求。
滑动窗口计数器

算法概念如下：

• 将时间划分为多个区间；
  
• 在每个区间内每有一次请求就将计数器加一维持一个时间窗口，占据多个区间；
  
• 每经过一个区间的时间，则抛弃最老的一个区间，并纳入最新的一个区间；
  
• 如果当前窗口内区间的请求计数总和超过了限制数量，则本窗口内所有的请求都被丢弃。
  

滑动窗口计数器是通过将窗口再细分，并且按照时间滑动，这种算法避免了固定窗口计数器带来的双倍突发请求，但时间区间的精度越高，算法所需的空间容量就越大。
go-zero 实现

go-zero 实现的是固定窗口的方式，计算一段时间内对同一个资源的访问次数，如果超过指定的 limit，则拒绝访问。当然如果在一段时间内访问不同的资源，每一个资源访问量都不超过 limit，此种情况是不会拒绝的。
而在一个分布式系统中，存在多个微服务提供服务。所以当瞬间的流量同时访问同一个资源，如何让计数器在分布式系统中正常计数？
这里要解决的一个主要问题就是计算的原子性，保证多个计算都能得到正确结果。
通过以下两个方面来解决：

• 使用 redis 的 incrby 做资源访问计数
  
• 采用 lua script 做整个窗口计算，保证计算的原子性
  

接下来先看一下 lua script 的源码：

1. // core/limit/periodlimit.go
3. const periodScript = `local limit = tonumber(ARGV[1])
4. local window = tonumber(ARGV[2])
5. local current = redis.call("INCRBY", KEYS[1], 1)
6. if current == 1 then
7.     redis.call("expire", KEYS[1], window)
8. end
9. if current < limit then
10.     return 1
11. elseif current == limit then
12.     return 2
13. else
14.     return 0
15. end`

复制代码

主要就是使用 INCRBY 命令来实现，第一次请求需要给 key 加上一个过期时间，到达过期时间之后，key 过期被清楚，重新计数。
限流器初始化：

1. type (
2.     // PeriodOption defines the method to customize a PeriodLimit.
3.     PeriodOption func(l *PeriodLimit)
5.     // A PeriodLimit is used to limit requests during a period of time.
6.     PeriodLimit struct {
7.         period     int  // 窗口大小，单位 s
8.         quota      int  // 请求上限
9.         limitStore *redis.Redis
10.         keyPrefix  string   // key 前缀
11.         align      bool
12.     }
13. )
15. // NewPeriodLimit returns a PeriodLimit with given parameters.
16. func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string,
17.     opts ...PeriodOption) *PeriodLimit {
18.     limiter := &PeriodLimit{
19.         period:     period,
20.         quota:      quota,
21.         limitStore: limitStore,
22.         keyPrefix:  keyPrefix,
23.     }
25.     for _, opt := range opts {
26.         opt(limiter)
27.     }
29.     return limiter
30. }

复制代码

调用限流：

1. // key 就是需要被限制的资源标识
2. func (h *PeriodLimit) Take(key string) (int, error) {
3.     return h.TakeCtx(context.Background(), key)
4. }
6. // TakeCtx requests a permit with context, it returns the permit state.
7. func (h *PeriodLimit) TakeCtx(ctx context.Context, key string) (int, error) {
8.     resp, err := h.limitStore.EvalCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
9.         strconv.Itoa(h.quota),
10.         strconv.Itoa(h.calcExpireSeconds()),
11.     })
12.     if err != nil {
13.         return Unknown, err
14.     }
16.     code, ok := resp.(int64)
17.     if !ok {
18.         return Unknown, ErrUnknownCode
19.     }
21.     switch code {
22.     case internalOverQuota: // 超过上限
23.         return OverQuota, nil
24.     case internalAllowed:   // 未超过，允许访问
25.         return Allowed, nil
26.     case internalHitQuota:  // 正好达到限流上限
27.         return HitQuota, nil
28.     default:
29.         return Unknown, ErrUnknownCode
30.     }
31. }

复制代码

上文已经介绍了，固定时间窗口会有临界突发问题，并不是那么严谨，下篇文章我们来介绍令牌桶限流。
以上就是本文的全部内容，如果觉得还不错的话欢迎点赞，转发和关注，感谢支持。
参考文章：

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

推荐阅读：

• go-zero 是如何做路由管理的？
  

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