Golang官方限流器time/rate的使用与实现详解

常见的限流算法

固定窗口计数器算法

固定窗口计数器算法将时间分为固定大小的窗口，例如1秒。在每个窗口中，服务会记录它接收到的请求数。如果在一个窗口中的请求数超过了预先设定的阈值，那么新的请求将被拒绝，直到进入下一个窗口。

这种算法简单易实现，但可能会导致窗口边界附近的请求突发。例如，如果窗口大小为1秒，阈值为100，那么在1秒的边界处，服务可能会在短时间内处理200个请求。

滑动窗口计数器算法

滑动窗口计数器算法试图解决固定窗口计数器算法中的请求突发问题。它将窗口分成更小的子窗口，例如将1秒分为10个100毫秒的子窗口。每次接收到请求时，服务会更新当前子窗口的计数器。服务会检查过去的N个子窗口的计数器之和，如果这个和超过阈值，那么新的请求将被拒绝。

这种算法可以更好地平滑请求流量，但实现起来相对复杂，因为需要跟踪多个子窗口的计数器。

令牌桶算法

令牌桶算法维护一个令牌桶，其中包含一定数量的令牌。令牌以恒定速率添加到桶中，直到达到桶的容量。每次接收到请求时，服务会尝试从桶中获取一个令牌。如果桶中有足够的令牌，请求被允许处理；如果没有足够的令牌，请求将被拒绝。

令牌桶算法允许短暂的请求突发，因为在低流量时期，令牌可以累积到桶的容量。这种算法在实践中表现良好，但实现起来相对复杂。

漏桶算法

漏桶算法使用一个队列模拟一个漏水的桶。请求作为水滴进入队列，以恒定速率从队列中移除并处理。如果队列已满，新的请求将被拒绝。

漏桶算法可以平滑请求流量，但它不能处理突发流量，因为请求处理速率是固定的。实现漏桶算法也相对复杂，因为需要在后台使用定时器或其他机制来以恒定速率处理队列中的请求。

time/rate

主要方法

• NewLimiter(limit Limit, burst int) *Limiter: 创建一个新的限流器，参数包括每秒允许的事件数量（limit）和令牌桶容量（burst）。

• (lim *Limiter) Allow() bool: 检查令牌桶中是否有可用的令牌。如果有可用令牌，则从桶中取走一个令牌并返回 true；否则返回 false。

• (lim *Limiter) AllowN(now time.Time, n int) bool: 与 Allow() 类似，但检查 n 个令牌是否可用。如果有足够的令牌，从桶中取走 n 个令牌并返回 true；否则返回 false。

• (lim *Limiter) Wait(ctx context.Context) error: 阻塞等待，直到有一个可用的令牌。如果在等待过程中 context 被取消或超时，将返回一个错误。

• (lim *Limiter) WaitN(ctx context.Context, n int) error: 阻塞等待，直到有 n 个可用的令牌。如果在等待过程中 context 被取消或超时，将返回一个错误。

• (lim *Limiter) Reserve() *Reservation: 返回一个预留令牌的 Reservation 对象。你可以根据需要等待预留令牌或取消预留。

• (lim *Limiter) ReserveN(now time.Time, n int) *Reservation: 类似于 Reserve()，但预留 n 个令牌。

各个方法的作用

• NewLimiter 用于创建一个新的限流器实例。

• Allow 和 AllowN 用于快速检查是否有足够的令牌可用，这些方法非阻塞。

• Wait 和 WaitN 用于阻塞等待直到有足够的令牌可用，这些方法会阻塞。

• Reserve 和 ReserveN 用于预留令牌，允许您根据需要等待预留令牌或取消预留。

time/rate 是如何实现限流的

time/rate 包基于令牌桶算法实现限流。限流器通过一个恒定速率（limit）向令牌桶添加令牌，直到桶的容量（burst）达到上限。每当处理一个请求时，限流器会尝试从令牌桶中取出一个或多个令牌。

Allow 和 AllowN 方法检查令牌桶中是否有足够的令牌。如果没有足够的令牌，这些方法会立即返回 false，表示应拒绝请求。Wait 和 WaitN 方法会阻塞等待，直到有足够的令牌可用。如果在等待过程中上下文（context）被取消或超时，这些方法会返回一个错误，表示请求被拒绝。Reserve 和 ReserveN 方法提供了更灵活的方式来预留令牌，您可以根据需要等待预留的令牌或取消预留。

通过这些方法，time/rate 限流器可以控制处理请求的速率，确保它不会超过设定的限制。通过调整令牌生成速率和令牌桶容量，您可以根据实际需求和系统负载来调整限流策略。

源码解析

令牌桶限流器的定义

在 rate.go 文件中，定义了 Limiter 结构体：

type Limiter struct {
   mu     sync.Mutex
   limit  Limit
   tokens float64
   // last 是上次令牌桶更新的时间
   last time.Time
   // 用于调整令牌桶更新时间的时钟
   clock Clock
   // 用于在 Wait 系列方法中进行休眠的定时器
   sleepFn func(time.Duration)
}

Limiter 结构体包含了一些关键属性，例如令牌生成速率（limit）、当前令牌数（tokens）和上次更新时间（last）。

令牌桶更新

time/rate 包中的核心函数之一是 reserveN，它负责预留 N 个令牌。在此过程中，令牌桶会根据时间更新。

func (lim *Limiter) reserveN(now time.Time, n int) *Reservation {
   lim.mu.Lock()
   defer lim.mu.Unlock()
   // 更新令牌桶
   now, tokens := lim.advance(now)
   // 计算需要的令牌数与当前可用令牌数之间的差值
   delta := float64(n) - tokens
   // 计算等待时间
   waitDuration := lim.limit.durationFromTokens(delta)
   // 更新令牌桶状态
   tokens -= float64(n)
   lim.last = now.Add(waitDuration)
   lim.tokens = tokens
   return &Reservation{
       ok:        true,
       lim:       lim,
       tokens:    n,
       timeToAct: now.Add(waitDuration),
   }
}

在 reserveN 函数中，首先调用 advance 函数来更新令牌桶：

func (lim *Limiter) advance(now time.Time) (time.Time, float64) {
   last := lim.last
   // 计算上次更新以来经过的时间
   elapsed := now.Sub(last)
   // 根据经过的时间计算生成的令牌数
   delta := elapsed.Seconds() * float64(lim.limit)
   // 更新令牌桶中的令牌数，但不超过令牌桶容量
   tokens := math.Min(lim.tokens+delta, float64(lim.limit.Burst()))
   return now, tokens
}

advance 函数根据时间更新令牌桶，计算从上次更新以来生成的令牌数量，并将新令牌添加到桶中，但不超过桶的容量。

令牌预留和等待

在 reserveN 函数中，首先计算需要的令牌数与当前可用令牌数之间的差值。然后根据差值计算等待时间。如果等待时间为正值，则表示需要等待一段时间

才能获得足够的令牌。最后，更新令牌桶状态，将所需令牌数从当前令牌数中减去。

reserveN 函数返回一个 Reservation 对象，其中包含预留的令牌数、等待时间等信息。Reservation 结构体定义如下：

type Reservation struct {
   ok        bool
   lim       *Limiter
   tokens    int
   timeToAct time.Time
}

Reservation 对象提供了一些方法，例如 Delay（返回需要等待的时间）和 Cancel（取消预留）。这些方法允许用户在需要时等待预留的令牌，或在不再需要令牌时取消预留。

公开 API

time/rate 包提供了一系列公开 API，例如 Allow, AllowN, Wait, WaitN, Reserve 和 ReserveN。这些方法都是基于 reserveN 函数的封装。例如，Allow 方法只需检查预留的等待时间是否为零：

func (lim *Limiter) Allow() bool {
   return lim.AllowN(time.Now(), 1)
}
func (lim *Limiter) AllowN(now time.Time, n int) bool {
   return lim.reserveN(now, n).Delay() == 0
}

类似地，Wait 和 WaitN 方法将阻塞等待，直到预留的等待时间过去：

func (lim *Limiter) Wait(ctx context.Context) error {
   return lim.WaitN(ctx, 1)
}
func (lim *Limiter) WaitN(ctx context.Context, n int) error {
   if n > lim.limit.Burst() {
       return fmt.Errorf("rate: Wait(n=％d) exceeds limiter's burst ％d", n, lim.limit.Burst())
   }
   r := lim.ReserveN(time.Now(), n)
   delay := r.DelayFrom(time.Now())
   if delay == 0 {
       return nil
   }
   t := lim.clock.AfterFunc(delay, r.Cancel)
   defer t.Stop()
   select {
   case <-ctx.Done():
       r.Cancel()
       return ctx.Err()
   case <-t.C:
       return nil
   }
}

总之，time/rate 包通过令牌桶算法实现了限流。它提供了一系列 API，允许用户在不同场景下灵活地控制请求速率。内部实现主要依赖于 reserveN 函数来更新令牌桶状态，并根据需要等待或预留令牌。

来源：https://segmentfault.com/a/1190000043736545