用 Go 写一个够用的令牌桶限流器

项目里有个对外接口需要按用户限流，每秒 5 次，允许短时间突发到 10 次。标准库扩展里的 golang.org/x/time/rate 完全够用，但这次想顺手搞清楚令牌桶到底怎么实现，就自己写了一版。

思路

令牌桶的核心只有两个数：桶里现在有多少令牌，以及上次更新的时间。不需要真的开一个 goroutine 每隔一段时间往桶里加令牌，每次请求进来时按时间差补上就行。

type Bucket struct {
    mu       sync.Mutex
    rate     float64   // 每秒补充的令牌数
    capacity float64   // 桶容量，也就是允许的突发量
    tokens   float64
    last     time.Time
}

func (b *Bucket) Allow() bool {
    b.mu.Lock()
    defer b.mu.Unlock()
    now := time.Now()
    b.tokens += now.Sub(b.last).Seconds() * b.rate
    if b.tokens > b.capacity {
        b.tokens = b.capacity
    }
    b.last = now
    if b.tokens >= 1 {
        b.tokens--
        return true
    }
    return false
}

几个容易写错的地方

初始令牌数。新建桶的时候 tokens 应该等于 capacity，否则服务刚启动时所有用户的第一个请求都会被拒绝。我第一版就是这么上线的，测试同学说"怎么第一次点总是失败"。

时间回拨。time.Now() 在 Go 里带单调时钟，Sub 用的是单调时钟读数，所以 NTP 调整不会让差值变成负数。但如果你把时间存到 Redis 里做分布式限流，就要自己处理这个问题了。

按用户分桶的内存。用 map[string]*Bucket 存每个用户的桶，用户多了内存会一直涨。我的做法是每个桶记一个最后访问时间，后台每分钟清理一次 10 分钟没访问过的桶。清理时满桶和被删掉的效果一样，所以不会影响限流的正确性。

压测结果

在一台 2 核的测试机上，单个桶 Allow() 大约 40ns，加上 map 查找和读写锁大约 120ns。对这个接口来说完全不是瓶颈。

最后还是换回了 x/time/rate，因为它的 Wait 和 Reserve 用起来更顺手。不过自己写一遍之后，再看它的源码就很快能看懂了。