限流是保证服务可用性和稳定的利器之一。go自带有限流器rate,我们来研读一下源码,源码较少,我们列出主要代码部分,注释说明:

// 限定速率
type Limit float64

// 常量Inf:无速率限制
const Inf = Limit(math.MaxFloat64)

// 限流器的结构体,本质就是一个令牌桶,
type Limiter struct {
mu sync.Mutex
limit Limit
burst int
tokens float64
// 记录上次 Limiter 被更新的时间
last time.Time
// lastEvent 记录速率受限制的时间
lastEvent time.Time
}


// new一个新的限流器 *Limiter,r是速率,b是允许突发的令牌数量
func NewLimiter(r Limit, b int) *Limiter {
return &Limiter{
limit: r,
burst: b,
}
}

// Reservation 保存了限流器延迟后发生的事件信息
type Reservation struct {
// 是否满足条件分配了令牌啊
ok bool
// 限流器
lim *Limiter
// 令牌数量
tokens int
// 满足令牌发放的时间
timeToAct time.Time
// 令牌发放速率
limit Limit
}


// 这是一个重要的内部函数,返回一个Reservation
func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
// 互斥锁,保证原子性
lim.mu.Lock()

// 判断速率是不是最大,是的话无需限流
if lim.limit == Inf {
lim.mu.Unlock()
return Reservation{
ok: true,
lim: lim,
tokens: n,
timeToAct: now,
}
}

// 调用 advance 方法,获取最新的时间、上一次取得令牌的时间、最新的token数量
now, last, tokens := lim.advance(now)

// 更新token数量
tokens -= float64(n)

// 计算等待时间,数量小于0表示令牌桶中木有可用的令牌了
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,
}

// 更新令牌桶数量和时间
if ok {
r.tokens = n
r.timeToAct = now.Add(waitDuration)
}

// 更新限流器最新取得令牌的时间、数量、事件
if ok {
lim.last = now
lim.tokens = tokens
lim.lastEvent = r.timeToAct
} else {
lim.last = last
}

lim.mu.Unlock()
return r
}

// 计算时间变化的 Limiter 的状态变化,得到最新的时间,最后一次取得令牌的时间和 token 令牌数量
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
last := lim.last
if now.Before(last) {
last = now
}

maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
elapsed := now.Sub(last)
if elapsed > maxElapsed {
elapsed = maxElapsed
}

delta := lim.limit.tokensFromDuration(elapsed)
tokens := lim.tokens + delta
if burst := float64(lim.burst); tokens > burst {
tokens = burst
}

return now, last, tokens
}


// 这里单独放出来,是提醒一下,go也有精度缺失问题,所以复杂计算要注意
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
sec := float64(d/time.Second) * float64(limit)
nsec := float64(d%time.Second) * float64(limit)
return sec + nsec/1e9
}

使用实例和demo可看:

golang 限流器