4. Go 实现令牌桶限流器：基于 time.Ticker 和 channel

令牌桶（Token Bucket）是限流领域最经典的算法之一。它通过一个固定容量的桶来存放令牌，以恒定速率向桶中添加令牌，请求到达时需从桶中获取一个令牌才能通过。这种设计既能平滑突发流量，又能保证长期的平均速率。

本节将使用 Go 语言的原生并发原语——time.Ticker 和 chan struct{}——实现一个生产可用的令牌桶限流器。代码不依赖第三方库，结构清晰，适合嵌入到微服务网关或 RPC 中间件中。

4.1 核心设计思路

令牌桶的实现可以拆解为两个并发角色：

  • 生产者：一个后台 goroutine 配合 time.Ticker，每隔固定时间间隔向桶中投放一个令牌。
  • 消费者：业务请求调用 Allow() 或 Wait() 方法，尝试从桶中取出一个令牌。

桶本身是一个带缓冲的 channel，缓冲大小即为桶的容量。channel 的收发语义天然符合令牌的“存入”与“取出”操作，且 goroutine-safe，无需额外加锁。

组件 类型 职责
桶 chan struct{} 存放令牌，缓冲长度 = 桶容量
生产者 time.Ticker + goroutine 按速率向 channel 发送空结构体
消费者 Allow() / Wait() 尝试从 channel 接收令牌
关闭控制 context.Context 优雅停止后台 goroutine
设计要点：使用 chan struct{} 而非 chan int 或 chan bool，因为空结构体不占用内存，且语义上仅表示“令牌存在”这一信号，无额外数据负载。

4.2 令牌桶结构体定义

// TokenBucket 令牌桶限流器
type TokenBucket struct {
    rate       time.Duration   // 令牌生成间隔
    capacity   int             // 桶容量
    tokens     chan struct{}   // 令牌 channel
    ctx        context.Context
    cancel     context.CancelFunc
}

4.3 构造函数与生产者启动

构造函数接收每秒期望的请求数（QPS）和桶容量。内部将 QPS 转换为令牌生成间隔，并启动后台 goroutine 持续向 channel 投放令牌。

// NewTokenBucket 创建一个令牌桶限流器
// qps: 每秒允许的请求数
// capacity: 桶容量，允许的突发大小
func NewTokenBucket(qps int, capacity int) *TokenBucket {
    if qps <= 0 || capacity <= 0 {
        panic("qps and capacity must be positive")
    }

    ctx, cancel := context.WithCancel(context.Background())

    tb := &TokenBucket{
        rate:     time.Second / time.Duration(qps),
        capacity: capacity,
        tokens:   make(chan struct{}, capacity),
        ctx:      ctx,
        cancel:   cancel,
    }

    // 预填充桶至满
    for i := 0; i < capacity; i++ {
        tb.tokens <- struct{}{}
    }

    // 启动生产者 goroutine
    go tb.produce()

    return tb
}

// produce 后台令牌生产者
func (tb *TokenBucket) produce() {
    ticker := time.NewTicker(tb.rate)
    defer ticker.Stop()

    for {
        select {
        case <-tb.ctx.Done():
            return
        case <-ticker.C:
            // 非阻塞发送：桶满时丢弃令牌
            select {
            case tb.tokens <- struct{}{}:
            default:
                // 桶已满，丢弃本次令牌
            }
        }
    }
}
关键细节：生产者使用双重 select 结构。外层 select 监听 ticker 和关闭信号；内层 select 向 channel 发送令牌时使用 default 分支处理桶满的情况，避免生产者阻塞。

4.4 消费者方法：Allow 与 Wait

限流器对外暴露两个核心方法：

  • Allow()：非阻塞尝试获取令牌，立即返回是否成功。
  • Wait()：阻塞直到获取到令牌，适用于必须等待的场景。
// Allow 非阻塞获取令牌
// 返回 true 表示获取成功，请求被允许
func (tb *TokenBucket) Allow() bool {
    select {
    case <-tb.tokens:
        return true
    default:
        return false
    }
}

// Wait 阻塞等待直到获取令牌
// 可通过 context 取消等待
func (tb *TokenBucket) Wait(ctx context.Context) error {
    select {
    case <-tb.tokens:
        return nil
    case <-ctx.Done():
        return ctx.Err()
    }
}

4.5 关闭与资源清理

限流器必须提供关闭方法，以停止后台 goroutine 并释放 ticker 资源。

// Stop 停止令牌桶，释放后台资源
func (tb *TokenBucket) Stop() {
    tb.cancel()
}

4.6 令牌桶工作流程

下图展示了令牌桶从初始化到请求限流的完整生命周期：

令牌桶限流器工作流程 令牌桶限流器工作流程 初始化桶（预填充） 生产者 goroutine time.Ticker 定时投放令牌 令牌桶 chan struct{} 缓冲大小 = capacity Allow() 非阻塞 立即返回 true / false Wait() 阻塞 等待直到获取令牌 请求通过 / 限流拒绝

4.7 完整使用示例

func main() {
    // 创建限流器：QPS=10，桶容量=5（允许突发5个请求）
    limiter := NewTokenBucket(10, 5)
    defer limiter.Stop()

    // 模拟20个并发请求
    var wg sync.WaitGroup
    for i := 0; i < 20; i++ {
        wg.Add(1)
        go func(id int) {
            defer wg.Done()
            if limiter.Allow() {
                fmt.Printf("[%d] 请求通过\n", id)
            } else {
                fmt.Printf("[%d] 请求被限流\n", id)
            }
        }(i)
    }
    wg.Wait()
}

4.8 性能与边界分析

场景 行为 说明
桶满时持续生产 令牌被丢弃 生产者 default 分支防止阻塞
桶空时大量请求 Allow 返回 false 非阻塞，请求快速失败
突发流量 桶中积攒的令牌被快速消耗 容量越大，突发容忍度越高
Stop 后继续调用 Allow 仍可消耗剩余令牌 生产者已停止，桶不再补充
注意事项：本实现基于 channel 和 goroutine，适用于大多数业务场景。但在极端高并发（百万 QPS 级别）下，channel 的调度开销可能成为瓶颈。此时建议使用原子操作 + 时间戳计算的纯用户态令牌桶实现。

4.9 与标准限流器对比

特性 本实现 golang.org/x/time/rate
算法 令牌桶（channel 实现） 令牌桶（浮点计算 + 原子操作）
并发安全 channel 天然安全 互斥锁 + 原子操作
突发支持 通过 capacity 控制 通过 Burst 方法控制
等待语义 Wait(ctx) 支持取消 Wait(ctx) 支持取消
性能 中等（goroutine 调度） 高（无 goroutine 开销）
依赖 纯标准库 需要引入 x/time 包

本实现的优势在于零依赖、代码可读性强、适合教学和中小规模项目。对于追求极致性能的生产环境，建议在理解原理后迁移到 x/time/rate 或自研原子操作版本。

← 上一章 📖 返回目录 下一章 →