Go 可用性 -二- 限流 1: 令牌桶原理及使用

序

從這篇文章開始會開始講一講限流該如何做，會結合 Go 進階訓練營中的內容、網上查閱的一些資料以及個人一些微小的經驗進行總結，這一部分預計總共會有 7 篇文章，今天我們來開始第一篇，令牌桶的原理及使用。

令牌桶

原理

如上圖 [5] 所示，令牌桶的大概原理是：

我們以 r/s 的速度向桶內放置令牌，桶的容量爲 b , 如果桶滿了令牌將會丟棄
當請求到達時，我們向桶內獲取令牌，如果令牌足夠，我們就通過轉發請求
如果桶內的令牌數量不夠，那麼這個請求會被緩存等待令牌足夠時轉發，或者是被直接丟棄掉

「由於桶的存在，所以令牌桶算法不僅可以限流還可以應對突發流量的情況」

舉個例子：假設我們桶的容量是 100，速度是 10 rps，那麼在我們桶滿的情況下，如果突然來 100 個請求是可以滿足的，但是後續的請求就會被限制到 10 rps

「存在下面兩種特殊情況」

如果桶的容量爲 0，那麼相當於禁止請求，因爲所有的令牌都被丟棄了
如果令牌放置速率爲無窮大，那麼相當於沒有限制

令牌桶最常見的實現就是 Go 官方的 golang.org/x/time/rate，接下來我們就看看這個庫如何使用吧

使用方法

方法一覽

// 限流器結構體
type Limiter
 // 構建一個限流器，r 是每秒放入的令牌數量，b 是桶的大小
    func NewLimiter(r Limit, b int) *Limiter

 // 分別返回 b 和 r 的值
    func (lim *Limiter) Burst() int
    func (lim *Limiter) Limit() Limit

 // token 消費方法
    func (lim *Limiter) Allow() bool
    func (lim *Limiter) AllowN(now time.Time, n int) bool
 func (lim *Limiter) Reserve() *Reservation
    func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation
    func (lim *Limiter) Wait(ctx context.Context) (err error)
    func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

 // 動態流控
    func (lim *Limiter) SetBurst(newBurst int)
    func (lim *Limiter) SetBurstAt(now time.Time, newBurst int)
    func (lim *Limiter) SetLimit(newLimit Limit)
    func (lim *Limiter) SetLimitAt(now time.Time, newLimit Limit)

初始化令牌桶

直接調用 NewLimiter(r Limit, b int) 即可， r 表示每秒產生 token 的速度， b 表示桶的大小

Token 消費

總共有三種 token 消費的方式，最常用的是使用 Wait 阻塞等待 **「Allow」

Allow 就是 AllowN(now,1) 的別名， AllowN 表示截止到 now 這個時間點，是否存在 n 個 token，如果存在那麼就返回 true 反之返回 false，如果我們限流比較嚴格，沒有資源就直接丟棄可以使用這個方法。

func (lim *Limiter) Allow() bool
func (lim *Limiter) AllowN(now time.Time, n int) bool

「Reserve」

同理 Reserve 也是 ReserveN(now, 1) 的別名， ReserveN 其實和 AllowN 類似，表示截止到 now 這個時間點，是否存在 n 個 token，只是 AllowN 直接返回 true or false，但是 ReserveN 返回一個 Reservation 對象

func (lim *Limiter) Reserve() *Reservation
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation

Reservation 有 5 個方法，通過調用 OK 我們可以知道是否通過等待可以獲取到 N 個 token，如果可以通過 Delay 方法我們可以得知需要等待的時間，如果我們不想等了可以調用 Cancel 方法歸還 token

type Reservation
    func (r *Reservation) Cancel()
    func (r *Reservation) CancelAt(now time.Time)
    func (r *Reservation) Delay() time.Duration
    func (r *Reservation) DelayFrom(now time.Time) time.Duration
    func (r *Reservation) OK() bool

「Wait」

Wait 是最常用的， Wait 是 WaitN(ctx, 1) 的別名， WaitN(ctx, n) 表示如果存在 n 個令牌就直接轉發，不存在我們就等，等待存在爲止，傳入的 ctx 的 Deadline 就是等待的 Deadline

func (lim *Limiter) Wait(ctx context.Context) (err error)
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error)

動態流控

通過調用 SetBurst 和 SetLimit 可以動態的設置桶的大小和 token 生產速率，其中 SetBurstAt 和 SetLimitAt 會將傳入的時間 now 設置爲流控最後的更新時間

func (lim *Limiter) SetBurst(newBurst int)
func (lim *Limiter) SetBurstAt(now time.Time, newBurst int)
func (lim *Limiter) SetLimit(newLimit Limit)
func (lim *Limiter) SetLimitAt(now time.Time, newLimit Limit)

案例: 10 行代碼實現基於 ip 的 gin 限流中間件

正好上週我們有個類似的需求，我們就來簡單實現一下，其實主要就是使用了 sync.map 來爲每一個 ip 創建一個 limiter，當然這個 key 也可以是其他的值，例如用戶名等

func NewLimiter(r rate.Limit, b int, t time.Duration) gin.HandlerFunc {
 limiters := &sync.Map{}

 return func(c *gin.Context) {
  // 獲取限速器
  // key 除了 ip 之外也可以是其他的，例如 header，user name 等
  key := c.ClientIP()
  l, _ := limiters.LoadOrStore(key, rate.NewLimiter(r, b))

  // 這裏注意不要直接使用 gin 的 context 默認是沒有超時時間的
  ctx, cancel := context.WithTimeout(c, t)
  defer cancel()

  if err := l.(*rate.Limiter).Wait(ctx); err != nil {
   // 這裏先不處理日誌了，如果返回錯誤就直接 429
   c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": err})
  }
  c.Next()
 }
}

使用的時候只需要 use 一下中間件就可以了

func main() {
 e := gin.Default()
 // 新建一個限速器，允許突發 10 個併發，限速 3rps，超過 500ms 就不再等待
 e.Use(NewLimiter(3, 10, 500*time.Millisecond))
 e.GET("ping", func(c *gin.Context) {
  c.String(http.StatusOK, "pong")
 })
 e.Run(":8080")
}

我們使用 go-stress-testing 來壓測一下，20 個併發

 ~/gopath/bin/go-stress-testing-mac -c 20 -n 1 -u http://127.0.0.1:8080/ping

開始啓動  併發數:20 請求數:1 請求參數:

─────┬───────┬───────┬───────┬────────┬────────┬────────┬────────┬────────┬────────┬────────
  耗時│  併發數│  成功數│  失敗數│   qps  │ 最長耗時│  最短耗時│ 平均耗時│  下載字節│ 字節每秒│ 錯誤碼
─────┼───────┼───────┼───────┼────────┼────────┼────────┼────────┼────────┼────────┼────────
   1s│     20│     11│      9│   63.79│  438.48│   45.37│  313.53│     152│     259│200:11;429:9


*************************  結果 stat  ****************************
處理協程數量: 20
請求總數（併發數*請求數 -c * -n）: 20 總請求時間: 0.586 秒 successNum: 11 failureNum: 9
*************************  結果 end   ****************************

可以發現總共成功了 11 個請求，失敗了 9 個，這是因爲我們桶的大小是 10 所以前 10 個請求都很快就結束了，第 11 個請求等待 333.3 ms 就可以完成，小於超時時間 500ms，所以可以放行，但是後面的請求確是等不了了，所以就都失敗了，並且可以看到最後一個成功的請求的耗時爲 336.83591ms 而其他的請求耗時都很短

[GIN-debug] Listening and serving HTTP on :8080
[GIN] 2021/03/29 - 13:15:55 | 200 |     1.48104ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |    1.107689ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |    1.746222ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |      866.35µs |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |    1.870403ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |    2.231912ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |    1.832506ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |     613.741µs |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |    1.454753ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |     1.37802ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |    1.428062ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |      40.782µs |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |    1.046146ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |      1.7624ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 429 |    1.803124ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |       41.67µs |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |     1.42315ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |    1.371483ms |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |     731.091µs |       127.0.0.1 | GET      "/ping"
[GIN] 2021/03/29 - 13:15:55 | 200 |   336.83591ms |       127.0.0.1 | GET      "/ping"

總結

這篇主要介紹了一下令牌桶的實現原理，以及 Go 官方實現如何使用，最後講了一個使用案例，下一篇我們再來學習一下源碼是怎麼實現的

參考文獻

Go 進階訓練營 - 極客時間
Token bucket - Wikipedia
令牌桶算法_百度百科 (baidu.com)
限流的概念，算法，分佈式限流以及微服務架構下限流的難點 - 知乎 (zhihu.com)
令牌桶工作原理 - 知乎 (zhihu.com)
golang.org/x/time/rate

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/xDk54SE-JPOsokFg12OKfA

序

令牌桶

原理