go语言后端开发学习(七)——如何在gin框架中集成限流中间件

一.什么是限流

限流又称为流量控制（流控），通常是指限制到达系统的并发请求数。
我们生活中也会经常遇到限流的场景，比如：某景区限制每日进入景区的游客数量为8万人；沙河地铁站早高峰通过站外排队逐一放行的方式限制同一时间进入车站的旅客数量等。
限流虽然会影响部分用户的使用体验，但是却能在一定程度上报障系统的稳定性，不至于崩溃（大家都没了用户体验）。
而互联网上类似需要限流的业务场景也有很多，比如电商系统的秒杀、微博上突发热点新闻、双十一购物节、12306抢票等等。这些场景下的用户请求量通常会激增，远远超过平时正常的请求量，此时如果不加任何限制很容易就会将后端服务打垮，影响服务的稳定性。
此外，一些厂商公开的API服务通常也会限制用户的请求次数，比如百度地图开放平台等会根据用户的付费情况来限制用户的请求数等。

二.常见的限流算法

2.1 漏桶算法

2.1.1 漏桶算法的原理

漏桶法的原理比较简单，假设我们有一个水桶按固定的速率向下方滴落一滴水，无论有多少请求，请求的速率有多大，都按照固定的速率流出，对应到系统中就是按照固定的速率处理请求。原理图如下：
在这里插入图片描述

漏桶法的关键点在于漏桶始终按照固定的速率运行，但是它并不能很好的处理有大量突发请求的场景，毕竟在某些场景下我们可能需要提高系统的处理效率，而不是一味的按照固定速率处理请求。

关于漏桶算法,在开发中我们可以使用三方的开源框架，uber团队有一个开源的github.com/uber-go/ratelimit库，下面网站是由漏桶算法集成以下如何简单的在gin中集成一个由漏桶算法实现的限流中间(这里我用的是我尝试自己编写的一个漏桶算法代码，大家可以选择自己写也可以选择使用上面的开源框架)：

//开源框架版
package main

import (
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/ratelimit"
)

func pong(c *gin.Context) {
	c.JSON(200, gin.H{
		"code":    200,
		"message": "pong",
	})

}

func LimitHandler() gin.HandlerFunc {
	return func(c *gin.Context) {
		r := ratelimit.New(1)  //每秒1个请求
		//每次滴水允许通过的请求数量
		// 限流
		if r.Take().Sub(time.Now()) > 0 {
			c.JSON(200, gin.H{
				"code":    429,
				"message": "Server busy",
			})
			c.Abort()
		}
	}
}

func main() {
	r := gin.Default()
	r.Use(LimitHandler())
	{
		r.GET("/ping", pong)
	}
	r.Run(":8080")
}

//自己实现版

//main.go
package main

import (
	"awesomeProject1/limit"
	"time"

	"github.com/gin-gonic/gin"
	"go.uber.org/ratelimit"
)

func pong(c *gin.Context) {
	c.JSON(200, gin.H{
		"code":    200,
		"message": "pong",
	})

}

func main() {
	r := gin.Default()
	r.Use(limit.LimitMiddleware())
	{
	}
	r.Run(":8080")
}

//limit.go
package limit

import (
	"sync"
	"time"

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

type Bucket struct {
	sync.Mutex
	lastAccess  time.Time
	requests    int64         // 当前已经接收请求次数
	MaxRequests int64         // 最大可接受请求次数
	interval    time.Duration // 时间间隔
}

func NewBucket(maxRequests int64, interval time.Duration) *Bucket {
	return &Bucket{
		lastAccess:  time.Now(),
		requests:    0,
		MaxRequests: maxRequests,
		interval:    interval,
	}
}

func (b *Bucket) Allow() bool {
	b.Lock() //加锁
	defer b.Unlock()
	now := time.Now()
	if now.Sub(b.lastAccess) > b.interval {
		b.requests = 0
		b.lastAccess = now
	}
	if b.requests < b.MaxRequests {
		b.requests++
		return true
	}
	return false
}

func LimitMiddleware() gin.HandlerFunc {
	bucket := NewBucket(1, 10*time.Second)
	return func(c *gin.Context) {
		if !bucket.Allow() {
			c.JSON(200, gin.H{
				"code":    429,
				"message": "Server busy",
			})
			c.Abort()
		}
		c.Next()
	}
}

2.2 令牌桶算法

2.2 令牌桶算法的原理

令牌桶其实和漏桶的原理类似，令牌桶会按固定的速率往桶里放入令牌，并且只要能从桶里取出令牌就能通过，令牌桶支持突发流量的快速处理。原理图如下：
在这里插入图片描述
当我们在令牌桶里面取不到令牌时我们就会选择拒绝该次请求。

2.2.2 基于令牌桶实现的限流中间件

和上面的漏桶限流一样，这里有关令牌桶的限流博主还是给出两个版本，一个是开源第三方库，同时博主也会写一个自己实现的限流中间件供大家参考：

首先是第三方库，这里我们可以考虑使用github.com/juju/ratelimit这一第三方库，下面我们来看一下如何基于这一第三方库封装出外面的限流中间件：

//     filepath:/limit/limiter
package limit

import (
	"time"

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

// 考虑到我们可能会对不同的请求做不同的限流，因此需要一个通用的实现接口
type LimiterInterface interface {
	Key(c *gin.Context) string                              //基于context实现获取对应限流器键值对
	GetBucket(key string) (*ratelimit.Bucket, bool)         // 获取限流器
	AddBuckets(rules ...LimiterBucketRule) LimiterInterface // 添加限流器规则
}

type Limiter struct{  // 限流器(用来记录不同接口对应的不同限流策略)
	LimiterBuckets map[string]*ratelimit.Bucket 
}

type LimiterBucketRule struct {  // 限流器规则
	Key          string        //
	FillInterval time.Duration // 时间间隔
	Capacity     int64         // 容量
	Quantum      int64         // 每次放置的令牌量
}


//  接口的具体实现   filepath:/limit/method_limiter.go
package limit

import (
	"strings"

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

type MethodLimiter struct {
	*Limiter
}

func NewMethodLimiter() LimiterInterface {
	l := &Limiter{LimiterBuckets: make(map[string]*ratelimit.Bucket)}
	return &MethodLimiter{Limiter: l}
}

func (l MethodLimiter) Key(c *gin.Context) string {
	url := c.Request.RequestURI
	index := strings.Index(url, "?")
	if index != -1 {
		url = url[:index]
	}
	return url
}

func (l MethodLimiter) GetBucket(key string) (*ratelimit.Bucket, bool) {
	bucket, ok := l.LimiterBuckets[key]
	return bucket, ok
}

func (l MethodLimiter) AddBuckets(rules ...LimiterBucketRule) LimiterInterface {
	for _, rule := range rules {
		if bucket, ok := l.LimiterBuckets[rule.Key]; !ok {
			bucket = ratelimit.NewBucketWithQuantum(rule.FillInterval, rule.Capacity, rule.Quantum)
			l.LimiterBuckets[rule.Key] = bucket
		}
	}
	return l
}

// 在gin框架中集成限流中间件        filepath: /middleware/limiter.go
package middleware

import (
	"awesomeProject1/limit"
	"fmt"

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

func LimitHandler(l limit.LimiterInterface) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := l.Key(c)
		if bucket, ok := l.GetBucket(key); ok {
			count := bucket.TakeAvailable(1)
			fmt.Println("key", key, "count", count)
			if count == 0 {
				c.JSON(429, gin.H{
					"code": 429,
					"msg":  "too many request",
				})
				c.Abort()
			}
		}
		c.Next()
	}
}

//测试样例 main.go
package main

import (
	"awesomeProject1/limit"
	"awesomeProject1/middleware"
	"time"

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

func pong(c *gin.Context) {
	c.JSON(200, gin.H{
		"code":    200,
		"message": "pong",
	})

}

func main() {
	r := gin.Default()
	limiter := limit.NewMethodLimiter()
	limiter.AddBuckets(
		limit.LimiterBucketRule{
			Key:          "/ping",
			FillInterval: 10 * time.Second,
			Capacity:     1,
			Quantum:      1,
		},
	)
	r.Use(middleware.LimitHandler(limiter))
	{
		r.GET("/ping", pong)
	}
	r.Run(":8080")
}