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) //登录


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

点赞1
点击评论0
收藏1
浏览 78
 

还没有评论,快来发表第一个评论吧

免责声明:凡在本网站出现的信息,均仅供参考。本网站将尽力确保所提供信息的准确性及可靠性,但不保证有关资料的准确性及可靠性,注册用户和一般页面游览者在使用前请进一步核实,并对任何自主决定的行为负责。本网站对有关资料所引致的错误、不确或遗漏,概不负任何法律责任(包括侵权责任、合同责任和其它责任)
*尊重作者,转载请注明出处!

创作内容

开启你的爱凌峰创作之旅

发布首篇内容,开通创作中心
快来成为爱凌峰创作者吧~

写文章

板块热门【Gin】