Go中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 个令牌。、
• SetLimit 和 SetBurst 方法：用于动态设置限制器的令牌生成速率和最大突发值。
• Limit 类型：用于表示每秒允许的最大事件数量。
• Limiter 结构体：包含限制器的状态，如令牌生成速率、最大突发值、令牌数量、最后更新时间等。
• Reservation 结构体：在调用 Reserve 或 Wait 方法后，返回一个表示事件是否允许发生以及需要等待多长时间的结构体。

time/rate 是如何实现限流的

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

1. 当事件发生时，用户通过调用 Allow、Reserve 或 Wait 方法尝试获取令牌。
2. 在获取令牌之前，Limiter 会根据当前时间和上次更新时间计算已生成的令牌数量。
3. 如果当前有足够的令牌可用，Limiter 会扣除相应的令牌数量，并允许事件发生。如果没有足够的令牌，Limiter 会根据调用的方法（Allow、Reserve 或 Wait）采取相应的行为。
4. 为了防止令牌桶过度填充，Limiter 会限制桶中令牌的最大数量为设置的突发值。

源码解析

令牌桶限流器的定义：

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

type Limiter struct {
    // 速率
    limit Limit
    // 突发值
    burst int
    
    // 锁
    mu     sync.Mutex
    // 令牌数
    tokens float64
    // last is the last time the limiter's tokens field was updated
    // 上次更新tokens最后一次发生时间
    last time.Time
    // lastEvent is the latest time of a rate-limited event (past or future)
    // 上次更新限速事件的最后一次发生时间。
    lastEvent time.Time
}

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

创建Limiter实例

NewLimiter函数用于创建一个新的Limiter实例：

func NewLimiter(r Limit, b int) *Limiter {
    return &Limiter{
        limit: r,
        burst: b,
        tokens: float64(b),
        last: time.Now(),
    }
}

NewLimiter接受速率和突发值作为参数，并初始化Limiter的字段。

令牌桶更新：

func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
    // 加锁以同步访问Limiter结构体
    lim.mu.Lock()

    // 如果速率是无限的，那么无需等待，直接返回Reservation
    if lim.limit == Inf {
        lim.mu.Unlock()
        return Reservation{
            ok:        true,
            lim:       lim,
            tokens:    n,
            timeToAct: now,
        }
    }

    // 调用advance方法更新令牌数量，获取当前时间、上次更新时间、令牌数量
    now, last, tokens := lim.advance(now)

    // 计算请求后剩余的令牌数量
    tokens -= float64(n)

    // 初始化等待时长
    var waitDuration time.Duration
    // 如果令牌数量不足，计算需要等待的时长
    if tokens < 0 {
        waitDuration = lim.limit.durationFromTokens(-tokens)
    }

    // 根据令牌数量和等待时长判断是否允许请求
    ok := n <= lim.burst && waitDuration <= maxFutureReserve

    // 准备Reservation实例
    r := Reservation{
        ok:    ok,
        lim:   lim,
        limit: lim.limit,
    }
    // 如果允许请求，设置Reservation的tokens和timeToAct
    if ok {
        r.tokens = n
        r.timeToAct = now.Add(waitDuration)
    }

    // 更新Limiter的状态
    if ok {
        lim.last = now
        lim.tokens = tokens
        lim.lastEvent = r.timeToAct
    } else {
        lim.last = last
    }

    // 解锁并返回Reservation实例
    lim.mu.Unlock()
    return r
}

reserveN函数首先加锁以同步访问Limiter结构体。然后，如果限制器的速率为无限，表示无需进行速率限制，直接返回一个Reservation实例，其中ok字段为true，表示请求被允许。

如果速率不为无限，函数会调用advance方法来更新令牌数量，并计算请求后剩余的令牌数量。接下来，如果令牌数量不足，函数会计算需要等待的时长。然后，根据令牌数量和等待时长判断请求是否被允许。

接着，函数会准备一个Reservation实例。如果请求被允许，函数会设置实例的tokens和timeToAct字段。最后，函数会根据请求是否被允许来

// advance 方法计算并返回速率限制器 lim 在经过一段时间后的状态。
// 请注意，lim 本身并不会被修改。
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
    // 保存 lim 的上一次事件时间
    last := lim.last

    // 如果当前时间 now 在上次事件时间 last 之前，将 last 设置为 now
    if now.Before(last) {
        last = now
    }

    // 计算当前剩余令牌数与允许的最大令牌数（即突发值）之间的差值所需的时间，以避免 delta 溢出
    maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
    // 计算自上次事件以来经过的时间
    elapsed := now.Sub(last)

    // 如果实际经过的时间 elapsed 大于 maxElapsed，则将 elapsed 设置为 maxElapsed
    if elapsed > maxElapsed {
        elapsed = maxElapsed
    }

    // 根据经过的时间 elapsed 计算新增的令牌数 delta
    delta := lim.limit.tokensFromDuration(elapsed)
    // 将新增的令牌数 delta 加到当前令牌数上
    tokens := lim.tokens + delta

    // 检查更新后的令牌数是否超过了突发值
    if burst := float64(lim.burst); tokens > burst {
        // 如果超过了突发值，将令牌数设置为突发值
        tokens = burst
    }

    // 返回更新后的当前时间、上次事件时间和令牌数
    return now, last,tokens
}

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

令牌预留和等待：

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

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

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

type Reservation struct {
    ok        bool
    lim       *Limiter
    tokens    int
    timeToAct time.Time
    // This is the Limit at reservation time, it can change later.
    limit Limit
}

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

公开 API：

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

// Allow is shorthand for AllowN(time.Now(), 1).
func (lim *Limiter) Allow() bool {
    return lim.AllowN(time.Now(), 1)
}

// AllowN reports whether n events may happen at time now.
// Use this method if you intend to drop / skip events that exceed the rate limit.
// Otherwise use Reserve or Wait.
func (lim *Limiter) AllowN(now time.Time, n int) bool {
    return lim.reserveN(now, n, 0).ok
}

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

func (lim *Limiter) Wait(ctx context.Context) error {
    return lim.WaitN(ctx, 1)
}

func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
    lim.mu.Lock()
    burst := lim.burst
    limit := lim.limit
    lim.mu.Unlock()

    if n > burst && limit != Inf {
        return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
    }
    // Check if ctx is already cancelled
    select {
    case <-ctx.Done():
        return ctx.Err()
    default:
    }
    // Determine wait limit
    now := time.Now()
    waitLimit := InfDuration
    if deadline, ok := ctx.Deadline(); ok {
        waitLimit = deadline.Sub(now)
    }
    // Reserve
    r := lim.reserveN(now, n, waitLimit)
    if !r.ok {
        return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
    }
    // Wait if necessary
    delay := r.DelayFrom(now)
    if delay == 0 {
        return nil
    }
    t := time.NewTimer(delay)
    defer t.Stop()
    select {
    case <-t.C:
        // We can proceed.
        return nil
    case <-ctx.Done():
        // Context was canceled before we could proceed.  Cancel the
        // reservation, which may permit other events to proceed sooner.
        r.Cancel()
        return ctx.Err()
    }
}

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