Golang（七）golang.org/x/time/rate 实现频率限制

1. 源码阅读

• 整个包实现原理基于令牌桶算法：随时间以 1/r 个令牌的速度向容积为 b 个令牌的桶中添加令牌，有请求就取走令牌，若令牌不足则不执行请求或者等待
• Allow 方法的调用链：lim.Allow() bool → lim.AllowN(time.Now(), 1) → lim.reserveN(now, n, 0).ok，因此 reserveN 方法的实现很关键

// Allow is shorthand for AllowN(time.Now(), 1).
func (lim *Limiter) Allow() bool {
    return lim.AllowN(time.Now(), 1)
}

// AllowN reports whether n events may happen at time now.
// Use this method if you intend to drop / skip events that exceed the rate limit.
// Otherwise use Reserve or Wait.
func (lim *Limiter) AllowN(now time.Time, n int) bool {
    return lim.reserveN(now, n, 0).ok
}

• reserveN 方法是线程安全的，通过互斥锁锁住判断操作：

lim.mu.Lock()
......
lim.mu.Unlock()
return r

• reserveN 方法，首先检查限制为 Inf，不需要判断直接返回 true：

if lim.limit == Inf {
    lim.mu.Unlock()
    return Reservation{
        ok:        true,
        lim:       lim,
        tokens:    n,
        timeToAct: now,
    }
}

• 时间和令牌的计算是根据单位时间内的限制数量和剩余令牌数所占比例计算的：

// durationFromTokens is a unit conversion function from the number of tokens to the duration
// of time it takes to accumulate them at a rate of limit tokens per second.
func (limit Limit) durationFromTokens(tokens float64) time.Duration {
    seconds := tokens / float64(limit)
    return time.Nanosecond * time.Duration(1e9*seconds)
}

// tokensFromDuration is a unit conversion function from a time duration to the number of tokens
// which could be accumulated during that duration at a rate of limit tokens per second.
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
    return d.Seconds() * float64(limit)
}

• 根据令牌数限制和当前存在的令牌数计算还有上次更新至今的时间差计算可以补充多少令牌：

// Avoid making delta overflow below when last is very old.
maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
elapsed := now.Sub(last)
if elapsed > maxElapsed {
    elapsed = maxElapsed
}

// Calculate the new number of tokens, due to time that passed.
delta := lim.limit.tokensFromDuration(elapsed)
tokens := lim.tokens + delta
if burst := float64(lim.burst); tokens > burst {
    tokens = burst
}

• 更新补充后的当前令牌数减去请求的令牌数，如果不足，根据不足的令牌数计算需要等待的时间：

// Calculate the remaining number of tokens resulting from the request.
tokens -= float64(n)
fmt.Printf("WangAo test: After sub n: tokens: %v
", tokens)

// Calculate the wait duration
var waitDuration time.Duration
if tokens < 0 {
    waitDuration = lim.limit.durationFromTokens(-tokens)
}

• 根据请求数量和等待时间判断是否允许请求：

// Decide result
ok := n <= lim.burst && waitDuration <= maxFutureReserve

3. 参考文献

• golang频率限制 golang.org/x/time/rate 使用说明

• 基于信号量和令牌桶算法的限流