行深善食-Go 进程的内存令牌桶限流

我目前先用的是进程内令牌桶(golang.org/x/time/rate),按 user_id/IP 维度限,单机场景足够。

// limiterEntry 每个 key 对应一个令牌桶 + 最后活跃时间
type limiterEntry struct {
	limiter  *rate.Limiter // 这个 key 对应的令牌桶
	lastSeen time.Time     // 这个 key 最后一次访问时间
}

var (
	limiterStore = sync.Map{} // key -> *limiterEntry
	cleanOnce    sync.Once
)

// getLimiter 按 key 拿(或新建)一个限流器
// r: 每秒允许多少个请求;b: 桶大小(允许突发多少)
func getLimiter(key string, r rate.Limit, b int) *rate.Limiter {
	cleanOnce.Do(startCleaner) // 启动后台清理协程(只启动一次)

	if v, ok := limiterStore.Load(key); ok {
		e := v.(*limiterEntry)
		e.lastSeen = time.Now()
		return e.limiter
	}

	e := &limiterEntry{
		// r 每秒往桶里放几个令牌(比如 `1` 就是每秒补 1 个);
		// b 桶容量(比如 `3` 就是最多允许一次性突发 3 个请求);
		limiter:  rate.NewLimiter(r, b),
		lastSeen: time.Now(),
	}
	actual, _ := limiterStore.LoadOrStore(key, e)
	return actual.(*limiterEntry).limiter
}

// RateLimit 通用限流中间件
// keyFunc 返回限流维度(user_id 或 client_ip)
func RateLimit(keyFunc func(*gin.Context) string, r rate.Limit, b int) gin.HandlerFunc {
	return func(c *gin.Context) {
		key := keyFunc(c)
		if key == "" {
			key = "unknown"
		}
		lim := getLimiter(key, r, b)
		if !lim.Allow() {
			c.AbortWithStatusJSON(429, &Response{
				Code: 429,
				Msg:  "请求过于频繁,请稍后再试",
				Data: nil,
			})
			return
		}
		c.Next()
	}
}

然后挂在路由里

// 已登录接口(饮食记录上报、建议提交):按 user_id 限
// 1 秒 5 个、突发 10
userWriteRateLimit := utils.RateLimit(utils.KeyByUserID, rate.Limit(5), 10)
authfoodGroup := healthGroup.Group("/food-group").Use(utils.JWTAuthMiddleware(), userWriteRateLimit)

用户瞬间发 10个请求,此时桶里有 10个令牌,全放行;第 11 个请求来了,此时桶空了,返回错误码429;之后每秒慢慢恢复 5 个令牌,每秒能过5 个;10 秒不访问,桶又攒满 10 个,允许一次突发 10 个。

用wrk压力测试

指想瞬间发 10 个请求,看前 3 个 200、后 7 个 429,用 bash 循环更直观:

for i in $(seq 1 30); do
    code=$(curl -s -o /dev/null -w "%{http_code}" \
        -H "Authorization: Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiLooYzmt7HohrPpo58iLCJleHAiOjE3OTczMDM4OTIsIm5iZiI6MTc4OTUyNzg5MiwiaWF0IjoxNzg5NTI3ODkyLCJ1c2VyX2lkIjo0NCwidXNlcm5hbWUiOiJ5YWxpaGFuIn0.TyiaFUAkL6mCghZuhZDV7Su9xu9aV4eji6CR-YsEEzc" \
        "http://192.168.31.109:30289/xs-json/health/food-group/user-composition?eating_pattern_id=3&estimated_calorie_needs_id=8")
    echo "请求 $i : $code"
done

结果如下:

请求 1 : 200
请求 2 : 200
请求 3 : 200
请求 4 : 200
请求 5 : 200
请求 6 : 200
请求 7 : 200
请求 8 : 200
请求 9 : 200
请求 10 : 200
请求 11 : 429
请求 12 : 429
请求 13 : 429
请求 14 : 429
请求 15 : 429
请求 16 : 429
请求 17 : 429
请求 18 : 429
请求 19 : 429
请求 20 : 429
请求 21 : 429
请求 22 : 429
请求 23 : 429
请求 24 : 429
请求 25 : 429
请求 26 : 429
请求 27 : 429
请求 28 : 429
请求 29 : 429
请求 30 : 429
[root@localhost ~]#