Gin 限流中间件

• 令牌桶

流量限制的手段有很多，最常见的：漏桶、令牌桶两种：

1. 漏桶是指我们有一个一直装满了水的桶，每过固定的一段时间即向外漏一滴水。如果你接到了这滴水，那么你就可以继续服务请求，如果没有接到，那么就需要等待下一滴水。
2. 令牌桶则是指匀速向桶中添加令牌，服务请求时需要从桶中获取令牌，令牌的数目可以按照需要消耗的资源进行相应的调整。如果没有令牌，可以选择等待，或者放弃。

这两种方法看起来很像，不过还是有区别的。漏桶流出的速率固定，而令牌桶只要在桶中有令牌，那就可以拿。也就是说令牌桶是允许一定程度的并发的，比如同一个时刻，有100个用户请求，只要令牌桶中有100个令牌，那么这100个请求全都会放过去。令牌桶在桶中没有令牌的情况下也会退化为漏桶模型，实际应用中令牌桶应用较为广泛。

令牌桶算法的原理是系统以恒定的速率产生令牌，然后把令牌放到令牌桶中，令牌桶有一个容量，当令牌桶满了的时候，再向其中放令牌，那么多余的令牌会被丢弃；当想要处理一个请求的时候，需要从令牌桶中取出一个令牌，如果此时令牌桶中没有令牌，那么则拒绝该请求。

对于令牌桶的Go语言实现，可以参照 github.com/juju/ratelimit 库。

• Gin 限流中间件

这里使用令牌桶作为限流策略，编写一个限流中间件如下：

package main

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/juju/ratelimit"
)

func RateLimitMiddleware(fillInterval time.Duration, cap, quantum int64) gin.HandlerFunc {
	bucket := ratelimit.NewBucketWithQuantum(fillInterval, cap, quantum)
	return func(c *gin.Context) {
		if bucket.TakeAvailable(1) < 1 {
			c.String(http.StatusForbidden, "rate limit...")
			c.Abort()
			return
		}
		c.Next()
	}
}

func main() {
	r := gin.Default()
	gin.ForceConsoleColor()
	r.Use(RateLimitMiddleware(time.Second, 100, 100))
	r.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "golang ~")
	})
	r.Run(":8080")
}

1. 令牌桶初始化后里面就有 100 个令牌
2. 每秒钟会产生 100 个令牌, 保证每秒最多有 100 个请求通过限流器, 也就是说 QPS 的上限是 100
3. 流量过大时能够启动限流, 在限流过程中, 仍然能让部分流量通过

github.com/juju/ratelimit 提供了几种不同特色的令牌桶填充方式：

• 默认的令牌桶，fillInterval 指每过多长时间向桶里放一个令牌，capacity 是桶的容量，超过桶容量的部分会被直接丢弃。桶初始是满的。

func NewBucket(fillInterval time.Duration, capacity int64) *Bucket

• 和普通的 NewBucket() 的区别是，每次向桶中放令牌时，是放 quantum 个令牌，而不是一个令牌。

func NewBucketWithQuantum(fillInterval time.Duration, capacity, quantum int64) *Bucket

• 这个就有点特殊了，会按照提供的比例，每秒钟填充令牌数。例如 capacity 是100，而 rate 是 0.1，那么每秒会填充10个令牌。

func NewBucketWithRate(rate float64, capacity int64) *Bucket

advanced-go-programming-book

Go 限流器实战系列 -- Token Bucket 令牌桶