聊聊 TokenBucket 限流器的基本原理及實現

限流器實現之 TokenBucket

大家好，我是「Go 學堂」的漁夫子。上篇文章我們講解了漏桶（LeakyBucket）的實現原理。本文我們介紹另外一種限流器 --- 令牌桶（TokenBucket）。

令牌桶（TokenBucket）簡介

令牌桶實現的基本思想

令牌桶，顧名思義，是一種通過讓請求被處理前先行獲取令牌，只有獲取到令牌的請求才能被放行處理的一種限流方式。令牌桶的實現包含兩個方面：

• 一方面是按固定的速率來產生令牌並存入桶中，如果令牌數量超過桶的最大容量則直接丟棄掉。
• 一方面當有請求時先從桶中獲取令牌，獲取到令牌後才能通過進行處理，否則被直接丟棄或等待獲取令牌。

令牌桶與漏桶（LeakyBucket）的區別

令牌桶與漏桶的區別在於漏桶控制的是請求被處理的速率。即當有請求的時候，先進入桶中進行排隊，按固定的速率流出被處理；而令牌桶控制的是令牌產生的速率。即當有請求的時候，先從令牌桶中獲取令牌，只要能獲取到令牌就能立即通過被處理，不限制請求被處理的速度，所以也就可以應對一定程度的突發流量。如圖所示二者的區別：

比如有 100 個請求同時進入。現在假設漏桶的速率是每 10ms 處理一個請求，那麼要處理完這 100 個請求需要 1 秒鐘，因為每處理完 1 個請求，都需要等待 10ms 才能處理下一個請求。

如果是令牌桶，假設令牌桶產生令牌的速率也是每 10ms 產生一個，那麼 1 秒鐘就是產生 100 個令牌。所以，一種極端的情況就是當這 100 個請求進入的時候，桶中正好有 100 個令牌，那麼這 100 個請求就能瞬間被處理掉。

golang 中的 time/rate 包

golang.org/x/time/rate 包就是基於令牌桶實現的。我們先來看下該包的使用，然後再分析該包的具體實現。

func main() {
  //構造限流器。第一個引數代表qps，即每秒鐘能夠產生多少令牌，第二個引數是指桶的最大容量，即最多能容下5個token
    limiter := NewLimiter(10, 5)

    for i := 0; i < 10; i++ {
        time.Sleep(time.Millisecond * 20)
        //Allow指去獲取令牌，如果沒獲取到，返回false
        if !limiter.Allow() {
            fmt.Printf("%d passed\n", i)
            continue
        }

        //說明請求通過Allow獲取到令牌了，繼續處理請求
        fmt.Println("%d droped\n" i)
        //todo 對請求的後續處理
    }
}

time/rate 實現原理

在簡介的部分我們提到，令牌桶需要按固定速率生成 Token。直觀的理解就是在令牌桶的實現中會有一個定時任務不斷的生成 Token。但在 Golang 的 time/rate 中的實現, 並沒有單獨維護一個定時任務，而是採用了 lazyload 的方式，直到每次有請求消費之前才根據時間差更新 Token 數目，同時通過計數的方式來計算當前桶中已有的 Token 數量。

Token 的生成和消耗

在開頭處我們提到，令牌桶是以固定的速率產生 Token，該速率就是我們在使用 NewLimiter 構造一個限流器時指定的第 1 個引數 limit，代表每秒鐘可以產生多少個 Token。

func NewLimiter(r Limit, b int) *Limiter {
    return &Limiter{
        limit: r,
        burst: b,
    }
}

那麼換算一下，就可以知道每生成一個 Token 的間隔時間 perToken = 1 秒/limit。假如我們指定的 limit 引數是 100，那麼 perToken=1 秒/100 = 10 毫秒。

上文提到，time/rate 包使用的是懶載入的方式生成的 Token。什麼是懶載入的方式呢？就是當有請求到來時，去桶中獲取令牌的同時先計算一下從上次生成令牌到現在的這段時間應該新增多少個令牌，把增量的令牌數先加到總的令牌資料上即可，後面被取走的令牌再從總數中減去即可。所以，我們在限流器中應該記錄下最近一次新增令牌的時間和令牌的總數，Limiter 的結構體會是如下這樣：

type Limiter struct {
    limit  Limit   //QPS，一秒鐘多少個token
    burst  int     //桶的容量，即最大能裝多少個令牌
    tokens float64 //當前的token數量
    last time.Time //last代表最近一次更新tokens的時間
}

好了，到這裡，我們有了生成一個 token 的時間間隔、最近一次更新 tokens 的時間、當前時間、當前的 token 數量四個屬性，就能很容易的計算每次有請求獲取令牌時，應該生成的令牌數量以及當前桶中總剩餘的令牌數了：

tokens += (當前時間 - 最近一次更新tokens的時間last) / 時間間隔

消耗的話，就是看當前的令牌總數是不是大於 0 就好了，如果大於 0，相當於該請求可以獲取令牌，從 tokens 中減 1，代表給該請求發放了一個令牌，該請求拿著令牌就能被通過進行處理了。

那到這裡是不是該演算法就結束了呢？並不是。那該 TokenBucket 是如何應對突發流量呢？

如何應對突發流量

所謂突發流量，就是在某個時刻的流量突然比平時的流量要高。光有突發流量還不夠，得系統能夠應對才行，即能夠正常處理，否則就不叫應對突發流量了。那令牌桶是如何應對突發流量的呢？就是通過令牌桶快取到的令牌來應對的，再加上令牌桶的最大容量約束，不會無限制的讓流量通過。 下面我們具體來看下應對突發流量的過程。

假設生成令牌的速率是每秒 100 個。而平常的請求平均值也就是每秒 80 個。也就是說每秒鐘令牌數能剩餘 20 個，那這剩餘的令牌就是用來應對突發流量的。例如在 10 秒後，就會有 200 個令牌剩餘，如果這個時候比平常多來 200 個請求，那麼令牌數也足以讓每個請求都領取到令牌從而被正常處理。

那麼，問題就又來了，如果在很長一段時間內，我們的系統請求數都很平穩，這樣我們就能積攢下很多剩餘的令牌，如果剩餘的令牌數很多，比如積攢了一千萬個了，突然來了一波流量，假設也是一千萬，按道理這一千萬個請求都能獲取到令牌，都能夠被處理。可問題是我們的計算機系統本身的資源缺不足以應付這麼多請求了，那該怎麼辦呢？

這個時候我們的令牌桶的最大容量屬性就該上場了，即 Limiter 結構體中的 burst 欄位。burst 欄位是限制該桶能夠儲存的最大令牌數，令牌積攢的數量超過該值後，就直接丟棄，不再進行積攢了。該欄位也代表我們的系統能夠應對的最大的突發流量數。 例如，一種極端的情況，在一段時間內，一個請求都沒有，但令牌會按照固定速率一直產生，這時令牌數達到了最大值 burst。突然有一波請求來了，假設是 burst+n 個，那這波流量中也就只有 burst 個請求能被正常處理，其他的就被拒絕了。因為我們只有 burst 個令牌。 所以，burst 值代表了我們系統應對突發流量的最大值。

數值溢位問題

我們在一開始講該演算法的實現時首先要計算從最後一次更新 tokens 數量到當前這段時間內產生的令牌數，以及令牌總的數量，一般的計算方式應該如下：

// elapsed表示最後一次更新tokens數量的時間到當前的時間差
elapsed := now.Sub(last)
// delta 具有數值溢位風險， 表示elapsed這段時間應該產生的令牌數量
delta := elapsed.Seconds() * float64(limit)

//tokens 表示當前總的令牌數量
tokens := lim.tokens + delta
if burst := float64(lim.burst); tokens > burst {
    tokens = burst
}

這裡有什麼問題呢？ 如果 last 值很小，那麼 elapsed 就會很大，而如果此時指定的 token 生成速率，即 limit 值也很大的話，那麼一個大值 乘以 一個大值，結果就很可能會溢位。

那該怎麼辦呢？ 我們知道，令牌桶有一個最大值 burst，如果超過這個 burst，那麼多餘的其實是沒用的。 因此，我們就可以先計算要填滿這個令牌桶最多需要多長時間 maxElapsed，如果時間差 now.Sub(last) 已經超過了該值，那麼說明令牌數就應該能達到最大值 burst 了。反之，說明 now.Sub(last) 是一個較小的值，繼續計算這段時間應該生成的令牌數即可，這樣就規避了大值相乘可能溢位的問題。 如下是 time/rate 包中的實現：

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
}

float64 精度問題

我們在上面 Limiter 的結構體中會注意到，tokens 的型別是 float64 的。你可能會問，難道令牌還有小數點，令牌數量不應該是整數嗎？ 是的，令牌數是有可能是小數的。為什麼呢？。

假設，我們指定的生成令牌的速率是每秒產生 965 個令牌，那麼每生成一個令牌的間隔是多少呢？大約是每 1.0362 毫秒產生一個令牌。那麼如果是在 100 毫秒這段時間會產生多少個令牌呢？大約 103.62 個令牌。

好了，既然是 float64，那麼在計算給定時間段內產生的 tokens 總數時就會有精度問題。我們來看看 time/rate 包的第一版的實現：

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
}

time.Duration 是 int64 的別名，代表納秒。分別求出秒的整數部分和小數部分，進行相乘後再相加，這樣可以得到最精確的精度。

總結

TokenBucket 是以固定的速率生成令牌，讓獲得令牌的請求才能通過被處理。令牌桶的限流方式可以應對一定的突發流量。在實現 TokenBucket 時需要注意在計算令牌總數時的數值溢位問題以及精度問題。