我目前先用的是进程内令牌桶(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 ~]#