前言：之前在Bluebell社區項目中，我們使用了開源的庫來實現令牌桶限流，這次我們試著使用Go來手撕實現下令牌桶算法。

1. 令牌桶算法

為了防止網絡擁塞，需要限制流出或者流入網絡的流量，使流量以比較均勻的速度向外發送。令牌桶算法就實現了這個功能，可控制發送到網絡上數據的數目，并允許突發數據的發送。

令牌桶算法是網絡流量整形和速率限制中最常使用的一種算法。大小固定的令牌桶可自行以恒定的速率源源不斷地產生令牌。如果令牌不被消耗，或者被消耗的速度小于產生的速度，令牌就會不斷地增多，直到把桶填滿。后面再產生的令牌就會從桶中溢出。最后桶中可以保存的最大令牌數永遠不會超過桶的大小。

傳送到令牌桶的數據包需要消耗令牌。不同大小的數據包，消耗的令牌數量不一樣。

令牌桶中的每一個令牌都代表一次請求或者一個字節。如果令牌桶中存在令牌，則允許發送流量；而如果令牌桶中不存在令牌，則不允許發送流量。因此，如果突發門限被合理地配置并且令牌桶中有足夠的令牌，那么流量就可以以峰值速率發送。

2. 調用第三方庫實現令牌桶

上次我們調用了封裝好的令牌桶算法，代碼如下。

回顧一下關鍵點，fillInterval time.Duration：表示令牌桶的填充間隔時間（例如，每秒填充一次）。cap int64：表示令牌桶的最大容量（即桶中最多可以存儲的令牌數量）。

它返回一個函數，該函數的簽名是 func(c *gin.Context)，這符合 gin 中間件的標準形式。也就是接收一個 *gin.Context 的參數，用于處理每個 HTTP 請求。

這里的 func(c *gin.Context) { ... } 是一個匿名函數。它捕獲 bucket 變量（閉包特性），因此可以在每次 HTTP 請求時使用同一個 bucket 實例。

func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(c *gin.Context) {bucket := ratelimit.NewBucket(fillInterval, cap)return func(c *gin.Context) {// 如果取不到令牌就中斷本次請求返回 rate limit...if bucket.TakeAvailable(1) == 0 {c.String(http.StatusOK, "rate limit...")c.Abort()return}// 取到令牌就放行c.Next()}
}

我們在路由組中使用Use調用中間件，代碼如下。

r.Use(Logger.GinLogger(), Logger.GinRecovery(stack: true), middlewares.RateLimitMiddleware(2*time.Second, cap: 1))

Logger.GinLogger()：添加了一個日志記錄中間件，用于記錄 HTTP 請求的日志信息。

Logger.GinRecovery(stack: true)：添加了一個錯誤恢復中間件，用于捕獲和處理在請求處理過程中發生的任何 panic，并可選地記錄堆棧跟蹤。

middlewares.RateLimitMiddleware(2*time.Second, cap: 1)：添加了速率限制中間件，配置為每兩秒鐘向令牌桶中添加一個令牌，桶的容量為 1。這意味著在任何給定的兩秒時間內，最多只能處理一個請求。

3. 手撕令牌桶

直接上代碼吧，邏輯比較簡單。咱們來看看總體代碼，再進行部分講解。

package awesomeProjectimport ("sync""time"
)// 定義令牌桶結構
type tokenBucket struct {limitRate int           // 限制頻率，即每分鐘加入多少個令牌tokenChan chan struct{} // 令牌通道，可以理解為桶cap       int           // 令牌桶的容量muLock    *sync.Mutex   // 令牌桶鎖，保證線程安全stop      bool          // 停止標記，結束令牌桶
}// NewTokenBucket 創建令牌桶
func NewTokenBucket(limitRate, cap int) *tokenBucket {if cap < 1 {panic("token bucket cap must be large 1")}return &tokenBucket{tokenChan: make(chan struct{}, cap),limitRate: limitRate,muLock:    new(sync.Mutex),cap:       cap,}
}// Start 開啟令牌桶
func (b *tokenBucket) Start() {go b.produce()
}// 生產令牌
func (b *tokenBucket) produce() {for {b.muLock.Lock()if b.stop {close(b.tokenChan)b.muLock.Unlock()return}b.tokenChan <- struct{}d := time.Minute / time.Duration(b.limitRate)b.muLock.Unlock()time.Sleep(d)}
}// Consume 消費令牌
func (b *tokenBucket) Consume() {<-b.tokenChan
}// Stop 停止令牌桶
func (b *tokenBucket) Stop() {b.muLock.Lock()defer b.muLock.Unlock()b.stop = true
}

---

func NewTokenBucket(limitRate, cap int) *tokenBucket {if cap < 1 {panic("token bucket cap must be large 1")}return &tokenBucket{tokenChan: make(chan struct{}, cap),limitRate: limitRate,muLock:    new(sync.Mutex),cap:       cap,}
}

tokenChan chan struct{}：tokenChan 是一個類型為 chan struct{} 的通道，用于存儲令牌。

make(chan struct{}, cap) 創建了一個容量為 cap 的緩沖通道。這里 struct{} 表示通道中存儲的是空結構體，僅用于占位，因為我們關心的是通道中元素的數量，而不是元素的值。這個通道模擬了令牌桶中令牌的數量，通道中的每個空結構體代表桶中的一個令牌。

func (b *tokenBucket) produce() {for {b.muLock.Lock()if b.stop {close(b.tokenChan)b.muLock.Unlock()return}b.tokenChan <- struct{}d := time.Minute / time.Duration(b.limitRate)b.muLock.Unlock()time.Sleep(d)}
}

time.Duration 是 Go 語言中用于表示時間間隔的類型，可以用于 time 包中的各種時間計算。
time.Minute 是一個常量，表示一分鐘的時間間隔。b.limitRate 是一個整數，表示每分鐘生成的令牌數量。

---

這里的通道chan是struct{}類型的，所以可以是別的。這里我們用struc{}空結構體，只是因為空結構體占用的內存空間非常小，適合用作通道中的占位符。