限流又称为流量控制（流控），通常是指限制到达系统的并发请求数，常用的限流算法主要有漏洞和令牌桶。

令牌桶

令牌桶其实和漏桶的原理类似，令牌桶按固定的速率往桶里放入令牌，并且只要能从桶里取出令牌就能通过，令牌桶支持突发流量的快速处理。
每日一库之105：juju/ratelimit - 图1
对于从桶里取不到令牌的场景，我们可以选择等待也可以直接拒绝并返回。
对于令牌桶的Go语言实现，大家可以参照github.com/juju/ratelimit库。这个库支持多种令牌桶模式，并且使用起来也比较简单。

创建令牌桶的方法：

  1. // 创建指定填充速率和容量大小的令牌桶
  2. func NewBucket(fillInterval time.Duration, capacity int64) *Bucket
  3. // 创建指定填充速率、容量大小和每次填充的令牌数的令牌桶
  4. func NewBucketWithQuantum(fillInterval time.Duration, capacity, quantum int64) *Bucket
  5. // 创建填充速度为指定速率和容量大小的令牌桶
  6. // NewBucketWithRate(0.1, 200) 表示每秒填充20个令牌
  7. func NewBucketWithRate(rate float64, capacity int64) *Bucket

取出令牌的方法如下：

  1. // 取token（非阻塞）
  2. func (tb *Bucket) Take(count int64) time.Duration
  3. func (tb *Bucket) TakeAvailable(count int64) int64
  5. // 最多等maxWait时间取token
  6. func (tb *Bucket) TakeMaxDuration(count int64, maxWait time.Duration) (time.Duration, bool)
  8. // 取token（阻塞）
  9. func (tb *Bucket) Wait(count int64)
  10. func (tb *Bucket) WaitMaxDuration(count int64, maxWait time.Duration) bool

虽说是令牌桶，但是我们没有必要真的去生成令牌放到桶里，我们只需要每次来取令牌的时候计算一下，当前是否有足够的令牌就可以了，具体的计算方式可以总结为下面的公式：
当前令牌数 = 上一次剩余的令牌数 + (本次取令牌的时刻-上一次取令牌的时刻)/放置令牌的时间间隔 * 每次放置的令牌数

github.com/juju/ratelimit这个库中关于令牌数计算的源代码如下：

  1. func (tb *Bucket) currentTick(now time.Time) int64 {
  2.     return int64(now.Sub(tb.startTime) / tb.fillInterval)
  3. }
  1. func (tb *Bucket) adjustavailableTokens(tick int64) {
  2.     if tb.availableTokens >= tb.capacity {
  3.         return
  4.     }
  5.     tb.availableTokens += (tick - tb.latestTick) * tb.quantum
  6.     if tb.availableTokens > tb.capacity {
  7.         tb.availableTokens = tb.capacity
  8.     }
  9.     tb.latestTick = tick
  10.     return
  11. }

获取令牌的TakeAvailable()函数关键部分的源代码如下：

  1. func (tb *Bucket) takeAvailable(now time.Time, count int64) int64 {
  2.     if count <= 0 {
  3.         return 0
  4.     }
  5.     tb.adjustavailableTokens(tb.currentTick(now))
  6.     if tb.availableTokens <= 0 {
  7.         return 0
  8.     }
  9.     if count > tb.availableTokens {
  10.         count = tb.availableTokens
  11.     }
  12.     tb.availableTokens -= count
  13.     return count
  14. }

大家从代码中也可以看到其实令牌桶的实现并没有很复杂。

gin框架中使用限流中间件

在gin框架构建的项目中，我们可以将限流组件定义成中间件。
这里使用令牌桶作为限流策略，编写一个限流中间件如下：

  1. func RateLimitMiddleware(fillInterval time.Duration, cap int64) func(c *gin.Context) {
  2.     bucket := ratelimit.NewBucket(fillInterval, cap)
  3.     return func(c *gin.Context) {
  4.         // 如果取不到令牌就中断本次请求返回 rate limit...
  5.         if bucket.TakeAvailable(1) < 1 {
  6.             c.String(http.StatusOK, "rate limit...")
  7.             c.Abort()
  8.             return
  9.         }
  10.         c.Next()
  11.     }
  12. }

对于该限流中间件的注册位置，我们可以按照不同的限流策略将其注册到不同的位置，例如：

  1. 如果要对全站限流就可以注册成全局的中间件。
  2. 如果是某一组路由需要限流，那么就只需将该限流中间件注册到对应的路由组即可。