示例说明
令牌桶的回填不需要后台线程——每次取令牌时按距离上次的时间差惰性补齐并封顶容量，数学上等价于连续回填。互斥锁保护桶状态的读改写；TryTake 的语义是立即判定，适合网关场景快速拒绝而不是让请求排队。按调用方分桶防止重客户挤占全局配额。
技术要点
- 惰性回填免后台 goroutine
- TryTake 语义适合快速拒绝
- 按 key 分桶实现调用方隔离
代码解读
take 是核心：先补令牌再判断，不足返回 false；容量封顶防止闲置期无限囤积。KeyedLimiter 维护 key 到桶的映射，双锁结构——外层短锁只负责查找或建桶，内层桶锁保护状态，不同 key 的限流互不阻塞。
使用建议
生产环境可直接用 golang.org/x/time/rate 标准实现。全局口径需外置 Redis。锁粒度已尽量收窄，超高并发可换原子操作版本。
代码示例
package main
import (
"math"
"sync"
"time"
)
type TokenBucket struct {
mu sync.Mutex
tokens float64
capacity float64
refill float64 // tokens per second
last time.Time
}
func NewTokenBucket(capacity float64, refillPerSec float64) *TokenBucket {
return &TokenBucket{
tokens: capacity,
capacity: capacity,
refill: refillPerSec,
last: time.Now(),
}
}
func (b *TokenBucket) take(now time.Time, n float64) bool {
elapsed := now.Sub(b.last).Seconds()
b.tokens = math.Min(b.capacity, b.tokens+elapsed*b.refill)
b.last = now
if b.tokens >= n {
b.tokens -= n
return true
}
return false
}
func (b *TokenBucket) TryTake(n float64) bool {
b.mu.Lock()
defer b.mu.Unlock()
return b.take(time.Now(), n)
}
type KeyedLimiter struct {
mu sync.Mutex
buckets map[string]*TokenBucket
capacity float64
refill float64
}
func (k *KeyedLimiter) TryTake(key string, n float64) bool {
k.mu.Lock()
b, ok := k.buckets[key]
if !ok {
b = NewTokenBucket(k.capacity, k.refill)
k.buckets[key] = b
}
k.mu.Unlock()
return b.TryTake(n)
}