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)}
}