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\n", 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 使用說明

• 基於信號量和令牌桶算法的限流