令牌桶算法原理

令牌桶（Token Bucket）是限流领域最经典的算法之一，也是 Go 标准库 rate.Limiter 的底层实现。理解它的数学模型与源码细节，是掌握 Go 限流能力的关键一步。

4.1 数学模型

令牌桶的核心思想是：系统以恒定速率向桶中放入令牌，每个请求必须从桶中获取一个令牌才能被处理。桶有最大容量，超出部分令牌被丢弃。

令牌桶的数学模型可以用以下公式描述：

// 当前令牌数 = min(桶容量, 上次令牌数 + (当前时间 - 上次更新时间) * 放入速率)
tokens = min(b, tokens + (now - lastTime) * r)

// 请求消耗 n 个令牌，成功条件：
if tokens >= n {
    tokens -= n
    return true  // 允许请求
}
return false     // 拒绝请求

这个模型保证了两个关键特性：

• 平滑突发：桶中积累的令牌允许短时间内的突发流量，突发上限为桶容量 b。
• 长期限速：长期平均速率不会超过 r，因为令牌的补充速率是恒定的。

4.2 令牌桶状态机

下图展示了令牌桶在请求到达时的完整状态流转：

4.3 Go 标准库 rate.Limiter 源码分析

Go 标准库 golang.org/x/time/rate 包中的 Limiter 是令牌桶的工业级实现。核心结构体如下：

type Limiter struct {
    mu     sync.Mutex
    limit  Limit        // 放入速率 r（每秒令牌数）
    burst  int          // 桶容量 b
    tokens float64      // 当前令牌数
    last   time.Time    // 上次更新令牌数的时间
    lastEvent time.Time // 上次限流事件的时间（用于 Wait 系列方法）
}

4.4 核心方法 reserveN 源码解析

reserveN 是 Limiter 最核心的内部方法，所有公开方法（Allow、Reserve、Wait）最终都调用它。以下是简化后的核心逻辑：

func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
    lim.mu.Lock()
    defer lim.mu.Unlock()

    // 1. 计算时间差，补充令牌
    last := lim.last
    tokens := lim.tokens
    elapsed := now.Sub(last)
    delta := elapsed.Seconds() * float64(lim.limit)
    tokens += delta
    if tokens > float64(lim.burst) {
        tokens = float64(lim.burst)
    }

    // 2. 计算消耗 n 个令牌后的令牌数
    tokens -= float64(n)

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

    // 4. 判断是否允许
    ok := n <= lim.burst && waitDuration <= maxFutureReserve

    // 5. 更新状态
    if ok {
        lim.last = now
        lim.tokens = tokens
        lim.lastEvent = now.Add(waitDuration)
    } else {
        lim.last = now
    }

    return Reservation{
        ok:    ok,
        lim:   lim,
        tokens: n,
        timeToAct: now.Add(waitDuration),
    }
}

4.5 三种调用方式对比

4.6 令牌桶 vs 漏桶

4.7 实战：构建一个带超时的限流器

以下示例展示了如何使用 rate.Limiter 实现一个带超时控制的限流器：

package main

import (
    "context"
    "fmt"
    "time"
    "golang.org/x/time/rate"
)

func main() {
    // 每秒 10 个令牌，桶容量 20
    limiter := rate.NewLimiter(10, 20)

    ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
    defer cancel()

    for i := 0; i < 30; i++ {
        if err := limiter.Wait(ctx); err != nil {
            fmt.Printf("请求 %d 被限流: %v\n", i, err)
            continue
        }
        fmt.Printf("请求 %d 通过\n", i)
    }
}

输出示例：

请求 0 通过
请求 1 通过
...
请求 19 通过
请求 20 被限流: rate: Wait(n=1) would exceed context deadline
请求 21 被限流: rate: Wait(n=1) would exceed context deadline

前 20 个请求利用桶中的积累令牌瞬间通过，后续请求因令牌补充速度（10个/秒）跟不上请求速度，且上下文超时，被拒绝。