Golang编写自定义IP限流中间件

基于令牌桶的限流算法

  • r:每秒钟向桶内放入 r 个令牌,即每隔 1/r秒放一个令牌
  • b:桶的最大容量是 b,桶满后试图再放入的令牌会被丢弃掉
  • 当有人请求 n 个令牌时,如果桶中的令牌数小于n,则请求放阻塞或直接放弃,否则顺利从桶中取走 n 个令牌
    • 当 b > 1 时,任意 1/r 秒内最多可以取走 b 个令牌
      • 限制很短的一个瞬间的一个最高的并发量,b=最高的并发量,r=最短的时间段
    • 当 b = 1 时,每秒钟最多可被取走 r 个令牌
      • 限制每秒钟的最高的QPS:b=1,r=最高的QPS
      • 限制每分钟的最高请求量:把每分钟的请求量/60=转换成1秒钟,赋给 r 就可以了
      • 限制每5分钟、每10分钟,同理

实现高并发限流(使用golang官方限流器)

Go代码

源码地址: GitHub-golang版本

middleware/rateLimiterMiddleware.go

package middleware

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

var Limiter *rate.Limiter

// 定义一个中间件函数来进行限流
func RateLimiterMiddleware() gin.HandlerFunc {
    
    
	return func(c *gin.Context) {
    
    
		if !Limiter.AllowN(time.Now(), 1) {
    
    
			c.JSON(http.StatusTooManyRequests, gin.H{
    
    "message": "Rate limit exceeded"})
			// 设置休眠和业务时长一样,为了更好从日志出看出规则
			time.Sleep(50 * time.Millisecond)
			c.Abort()
			return
		}

		c.Next()
	}
}

main.go

func main() {
    
    
	r := gin.Default()
	
	// 创建一个限流器,每秒允许最多10个请求
	middleware.Limiter = rate.NewLimiter(rate.Limit(10), 1)

	// 使用限流中间件
	r.Use(middleware.RateLimiterMiddleware())

	r.GET("/api/resource", func(c *gin.Context) {
    
    
		time.Sleep(50 * time.Millisecond)
		c.JSON(http.StatusOK, gin.H{
    
    "message": "Resource accessed"})
	})

	r.Run(":8080")
}

测试记录

ab -t 1 -c 1 http://127.0.0.1:8080/api/resource

使用ab压力测试,并发量为1的(相当于串行),在1秒内不断发出请求(算下来,每个请求50ms,总共能发出20个请求)

结果预测:1秒内最多生成10个令牌,而总共有20个串行的请求,结果应该是1个成功(在50ms结束),1个失败(后50ms内还未有新的令牌生成),1个成功,1个失败。。。

结果输出(符合预期)

在这里插入图片描述
在这里插入图片描述

升级:根据每个IP地址进行限流

Go代码

源码地址: GitHub-golang版本

middleware/ipRateLimiterMiddleware.go

package middleware

import (
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)

var IPLimiter *IPRateLimiter

func NewIPRateLimiter() *IPRateLimiter {
    
    
	return &IPRateLimiter{
    
    
		limiter: make(map[string]*rate.Limiter),
	}
}

type IPRateLimiter struct {
    
    
	mu      sync.Mutex
	limiter map[string]*rate.Limiter
}

func (i *IPRateLimiter) GetLimiter(ip string) *rate.Limiter {
    
    
	i.mu.Lock()
	defer i.mu.Unlock()

	limiter, exists := i.limiter[ip]
	if !exists {
    
    
		limiter = rate.NewLimiter(2, 5) // 每秒2个请求,桶容量为5
		i.limiter[ip] = limiter
	}

	return limiter
}

// 定义一个中间件函数来进行限流
func IPRateLimiterMiddleware() gin.HandlerFunc {
    
    
	return func(c *gin.Context) {
    
    
		ip := c.ClientIP()

		limiter := IPLimiter.GetLimiter(ip)

		if !limiter.Allow() {
    
    
			c.JSON(http.StatusTooManyRequests, gin.H{
    
    "message": "Rate limit exceeded"})
			// 设置休眠和业务时长一样,为了更好从日志出看出规则
			time.Sleep(50 * time.Millisecond)
			c.Abort()
			return
		}

		c.Next()
	}
}

main.go

func main() {
    
    
	r := gin.Default()

	// 创建IP限流器
	middleware.IPLimiter = middleware.NewIPRateLimiter()

	// 使用限流中间件
	r.Use(middleware.IPRateLimiterMiddleware())

	r.GET("/api/resource", func(c *gin.Context) {
    
    
		time.Sleep(50 * time.Millisecond)
		c.JSON(http.StatusOK, gin.H{
    
    "message": "Resource accessed"})
	})

	r.Run(":8080")
}

测试记录

ab -t 1 -c 1 http://127.0.0.1:8080/api/resource

使用ab压力测试,并发量为1的(相当于串行),在1秒内不断发出请求(算下来,每个请求50ms,总共能发出20个请求)

结果预测:1秒内最多生成2个令牌,桶容量为5 代表在1/2秒内的最大并发量是5,总共有20个串行的请求,结果应该是先成功5个(桶容量全使用成功), 之后剩余成功的是1个,其余全部失败

  1. 当程序启动时,限流器初始化,桶是空的,没有令牌。
  2. 在每秒的前两次令牌生成中,每次生成2个令牌,并放入桶中。
  3. 在第3秒,生成的2个令牌中只有1个能够被放入桶中,因为桶已经满了。
  4. 当使用ab发出请求时,前5个请求依次拿到了令牌并成功,每个请求耗时50ms。此时总耗时为250ms。
  5. 第6个请求在300ms,第7个请求在350ms,第8个请求在400ms,第9个请求在450ms,第10个请求在500ms。
  6. 因为限流器的速率是每秒2个请求,也就是每500ms生成一个令牌。 第500ms时新的令牌生成,所以第6-10请求期间,是没有令牌的,所以都请求失败,而11个请求时有新令牌生成,所以能成功。
  7. 之后的请求以此类推,需要等待新令牌生成。

结果输出(符合预期)

在这里插入图片描述
在这里插入图片描述

猜你喜欢

转载自blog.csdn.net/trinityleo5/article/details/132889703