1. 前言
在上一篇文章 Go微服務: 令牌桶 當中簡單的介紹了令牌桶實現的原理，然後利用 /x/time/rate 這個庫 10 行程式碼寫了一個基於 ip 的 gin 限流中介軟體，那這個功能是怎麼實現的呢？接下來我們就從原始碼層面來了解一下這個庫的實現。這個實現很有意思，並沒有真正的使用一個定時器不斷的生成令牌，而是靠計算的方式來完成
2.rate/limt
在golang.org/x/time/rate庫中
使用限速器的時候我們需要呼叫 NewLimiter 方法，然後 Limiter 提供了三組限速的方法，這三組方法其實都是通過呼叫 reserveN 實現的 reserveN 返回一個 *Reservation 指標，先來看一下這兩個結構體。
2.1 Limiter
type Limiter struct {
// 互斥鎖
mu sync.Mutex
// 每秒產生 token 的速度, 其實是 float64 的一個別名
limit Limit
// 桶的大小
burst int
// 當前時間節點擁有的 tokens 數量
tokens float64
// 上次更新 token 的時間
last time.Time
// 上次限速的時間，這個時間可能是過去的某個時間也可能是將來的某個時間
lastEvent time.Time
}
2.2 Reservation
預定，表示預約某個時間的 token
type Reservation struct {
// 是否能預約上
ok bool
// limter
lim *Limiter
// 預約的 token 數量
tokens int
// token 實際使用的時間
timeToAct time.Time
// 儲存一下速率，因為 lim 的速率是可以被動態調整的，所以不能直接用
limit Limit
}
這個庫並沒有使用定時器來發放 token 而是用了 lazyload 的方式，等需要消費 token 的時候才通過時間去計算然後更新 token 的數量，下面我們先通過一個例子來看一下這個流程是怎麼跑的
如上圖所示，假設我們有一個限速器，它的 token 生成速度為 1，也就是一秒一個，桶的大小為 10，每個格子表示一秒的時間間隔
last表示上一次更新 token時還有 2 個token- 現在我們有一個請求竟來， 總共需要7個 token才能完成請求
now表示我現在進來的時間，距離last 已經過去了2s, 那麼現在就有4個token(每秒生成一個token)- 所以,如果需要 7 個 token 那麼也就還需要等待 3s 中才真的有 7 個，所以這就是 timeToAct 所在的時間節點
- 預約成功之後更新 last = now 、token = -3 因為 token 已經被預約出去了所以現在剩下的就是負數了
2.3 消費 token
總共有三種消費 token 的方法 AllowN, ReserveN, WaitN最終都是呼叫的reserveN 這個方法
// now: 需要消費 token 的時間點
// n: 需要多少個 token
// maxFutureReserve: 能夠等待的最長時間
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 可能會變
now, last, tokens := lim.advance(now)
// 發放 token 之後還剩多少
tokens -= float64(n)
// 根據 token 數量計算需要等待的時間
var waitDuration time.Duration
if tokens < 0 {
waitDuration = lim.limit.durationFromTokens(-tokens)
}
// 計算是否可以發放，如果需要的量比桶的容量還大肯定是不行的
// 然後就是看需要能否容忍需要等待的時間
ok := n <= lim.burst && waitDuration <= maxFutureReserve
// Prepare reservation
r := Reservation{
ok: ok,
lim: lim,
limit: lim.limit,
}
// 如果可以的話，就把 token 分配給預約者
if ok {
r.tokens = n
r.timeToAct = now.Add(waitDuration)
}
// 更新各個欄位的狀態
if ok {
lim.last = now
lim.tokens = tokens
lim.lastEvent = r.timeToAct
} else {
// 為什麼不 ok 也要更新 last 呢？因為 last 可能會改變
lim.last = last
}
lim.mu.Unlock()
return r
}
advance 方法用於計算 token 的數量
// now 是傳入的當前的時間點，返回的 newNow 其實就是傳入的引數，沒有任何改變
// newLast 是更新 token 的時間
// newTokens 是 token 的數量
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
// 如果當前時間比上次更新 token 的時間還要早，那麼就重置一下 last
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
}
// 計算這段時間生成的 token 數量，如果大於桶的容量，就取桶的容量
delta := lim.limit.tokensFromDuration(elapsed)
tokens := lim.tokens + delta
if burst := float64(lim.burst); tokens > burst {
tokens = burst
}
return now, last, tokens
}
這個比較有意思的是先去計算了時間的最大值，因為初始化的時候沒為 last 賦值，所以 now.Before(last) 出來的結果可能是一個很大的值，再去計算 tokens 數量很可能溢位
durationFromTokens 根據 tokens 的數量計算需要花費的時間
func (limit Limit) durationFromTokens(tokens float64) time.Duration {
seconds := tokens / float64(limit)
return time.Nanosecond * time.Duration(1e9*seconds)
}
tokensFromDuration根據時間計算 tokens 的數量
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
}
2.4 消費token的總結
消費 token 的邏輯就講完了，大概總結一下
- 需要消費的時候， 先去計算一下，從過去到現在可以生成多少個token
- 然後通過需要的 token 減去現在擁有的token數量，就得到了需要預約的token數量
- 再通過token數量 轉換成時間，就可以得到需要等待的時間長度，以及是否可以消費
- 然後再通過不同的消費方式進行消費
2.5 WaitN
// ctx 用於控制超時, n 是需要消費的 token 數量，如果 context 的 Deadline 早於要等待的時間就會直接返回失敗
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
lim.mu.Lock()
burst := lim.burst
limit := lim.limit
lim.mu.Unlock()
// 先看一下是不是已經超出消費極限了
if n > burst && limit != Inf {
return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, burst)
}
// 如果 ctx 已經結束了也不用等了
select {
case <-ctx.Done():
return ctx.Err()
default:
}
// 計算一下可以等待的時間
now := time.Now()
waitLimit := InfDuration
if deadline, ok := ctx.Deadline(); ok {
waitLimit = deadline.Sub(now)
}
// 呼叫 reserveN 得到預約資料
r := lim.reserveN(now, n, waitLimit)
// 如果不 ok 說明預約不到
if !r.ok {
return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
}
// 如果可以預約到，計算一下需要等多久
delay := r.DelayFrom(now)
if delay == 0 {
return nil
}
// 啟動一個 timer 進行定時
t := time.NewTimer(delay)
defer t.Stop()
select {
case <-t.C:
// We can proceed.
return nil
case <-ctx.Done():
// 如果 context 主動取消了，那麼之前預約的 token 數量需要歸還
r.Cancel()
return ctx.Err()
}
}
2.5 取消消費
WaitN 當中如果預約上了，但是 Context 取消了，會呼叫 CancelAt 歸還 tokens， 實現原理如下
func (r *Reservation) CancelAt(now time.Time) {
// 不 ok 說明沒有預約上，直接返回就行了
if !r.ok {
return
}
r.lim.mu.Lock()
defer r.lim.mu.Unlock()
// 如果沒有速率限制，或者沒有消費 token 或 token 已經被消費了，都不用還了
if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(now) {
return
}
// 計算需要還的 token 數量
// 這裡說是需要減去已經預支的 token 數量，但是我發現應該是個 bug，感覺這裡減重複了
restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
if restoreTokens <= 0 {
return
}
// 計算當前擁有的 tokens 數量
now, _, tokens := r.lim.advance(now)
// 當前擁有的加上需要歸還的就是現有的，但是不能大於桶的容量
tokens += restoreTokens
if burst := float64(r.lim.burst); tokens > burst {
tokens = burst
}
// 更新 tokens 數量
r.lim.last = now
r.lim.tokens = tokens
// 如果相等說明後面沒有新的 token 消費，所以將狀態重置到上一次
if r.timeToAct == r.lim.lastEvent {
prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
if !prevEvent.Before(now) {
r.lim.lastEvent = prevEvent
}
}
return
}
3. 存在的問題
除了上面提到的感覺 cancelAt 可能有一個 bug 外，雲神的部落格還提到了一個問題，就是如果我們 cancel 了的話，後面已經在等待的任務是不會重新調整的，舉個例子
func wait() {
l := rate.NewLimiter(10, 10)
t := time.Now()
l.ReserveN(t, 10)
var wg sync.WaitGroup
ctx, cancel := context.WithTimeout(context.TODO(), time.Hour)
defer cancel()
// 註釋掉下面這段就不會提前 cancel
wg.Add(1)
go func() {
defer wg.Done()
// 模擬出現問題, 200ms就取消了
time.Sleep(200 * time.Millisecond)
cancel()
}()
wg.Add(2)
go func() {
defer wg.Done()
// 如果要等，這個要等 1s 才能執行，但是我們的 ctx 200ms 就會取消
l.WaitN(ctx, 10)
fmt.Printf("[1] cost: %s\n", time.Since(t))
}()
time.Sleep(100 * time.Millisecond)
go func() {
defer wg.Done()
// 正常情況下，這個要等 1.2 s 才能執行，但是我們前面都取消了
// 這個是不是應該就只需要等 200ms 就執行了
ctx, cancel := context.WithTimeout(context.Background(), time.Hour)
defer cancel()
l.WaitN(ctx, 2)
fmt.Printf("[2] cost: %s\n", time.Since(t))
}()
wg.Wait()
}
先看一下不提前 cancel 的結果
[1] cost: 1.0002113s
[2] cost: 1.2007347s
再看看提前 cancel 的結果
[1] cost: 200.8268ms
[2] cost: 1.201066s
可以看到就是 1 有變化，從 1s -> 200ms 但是 2 一直都要等 1.2s
3.1 關於可能存在的bug
restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
在取消的時候，會減掉一個預約的時間，但是我發現這裡其實應該是重複減了一次
測試程式碼
func main() {
t0 := time.Now()
t1 := time.Now().Add(100 * time.Millisecond)
t2 := time.Now().Add(200 * time.Millisecond)
t3 := time.Now().Add(300 * time.Millisecond)
l := rate.NewLimiter(10, 20)
l.ReserveN(t0, 15) // 桶裡還剩 5 個 token
fmt.Printf("%+v\n", l)
r := l.ReserveN(t1, 10) // 桶還有 -4 個，
fmt.Printf("%+v\n", l)
// 註釋掉下面兩行，最後結果還剩 8 個 token
l.ReserveN(t2, 2) // 桶裡還有 -5 個
fmt.Printf("%+v\n", l)
r.CancelAt(t3)
fmt.Printf("%+v\n", l)
// 歸還之前借的，執行結果 桶裡還有 4 個
// 但是這裡不應該剩下 6 個麼，本來有 5 個，300ms 生成了 3 個，後面又預支出去 2 個
// 而且我發現如果我註釋掉預支兩個的程式碼，結果和我預期的一致，剩餘 8 個token
}
4. 總結
這一節主要是看了原始碼，但是其中還是有很多的細節沒有深入的去了解，後面再去細看吧