Gin框架系列教程（10）- gin IP限流+令牌桶限流

   我们经常有一些场景需要用到限流，比如，我们登录的时候会限制IP时间段内登录的次数，或者防止并发量突然增高时，服务器无法承受，保证了QPS的上限值。

限流的方法有很多，下面我们来实现一下gin IP限流+令牌桶限流，参考网上多种方法测试代码如下：

中间件：

package middleware

import (
	"admin/services/commonService/response"
	"fmt"
	"net/http"
	"sync"
	"time"

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

type RequestInfo struct {
	LastAccessTime time.Time // 上次访问时间
	RequestNum     int       // 请求计数
}

var (
	requestInfoMap = make(map[string]*RequestInfo) // IP到请求信息的映射
	mutex          = &sync.Mutex{}                 // 用于保护requestInfoMap的互斥锁
	maxRequests    = 2                             // 允许的最大请求数
	timeWindow     = 1 * time.Second               // 时间窗口
)

// IP限流器
func IpLimit(c *gin.Context) {
	ip := c.ClientIP()
	mutex.Lock()
	defer mutex.Unlock()
	// 检查IP是否在map中
	info, exists := requestInfoMap[ip]
	// 如果IP不存在，初始化并添加到map中
	if !exists {
		requestInfoMap[ip] = &RequestInfo{LastAccessTime: time.Now(), RequestNum: 1}
		return
	}
	// 如果IP存在，检查时间窗口
	if time.Since(info.LastAccessTime) > timeWindow {
		// 如果超过时间窗口，重置请求计数
		info.RequestNum = 1
		info.LastAccessTime = time.Now()
		return
	}
	info.RequestNum++ // 如果在时间窗口内，增加请求计数
	// 如果请求计数超过限制，禁止访问
	if info.RequestNum > maxRequests {
		errMsg := fmt.Errorf("too many requests")
		response.Error(c, http.StatusTooManyRequests, errMsg, "请求过于频繁，请稍后再试！")
		c.Abort()
		return
	}
	// 更新最后访问时间
	info.LastAccessTime = time.Now()
	c.Next()
}

// ratelimit限流器
func RateLimit(time time.Duration, originNum, pushNum int64) gin.HandlerFunc {
	bucket := ratelimit.NewBucketWithQuantum(time, originNum, pushNum)
	return func(c *gin.Context) {
		if bucket.TakeAvailable(1) < 1 {
			errMsg := fmt.Errorf("too many requests")
			response.Error(c, http.StatusTooManyRequests, errMsg, "请求过于频繁，请稍后再试！")
			c.Abort()
			return
		}
		c.Next()
	}
}

路由中使用中间件：

api.POST("/member/login", middleware.IpLimit, memberApi.MemberLogin) //登录
api.POST("/member/login", middleware.RateLimit(time.Second, 10, 10), memberApi.MemberLogin) //登录

结果：经过测试上面两种方式都能达到限流的目的。