diff --git a/init/ratelimit/ratelimit.go b/init/ratelimit/ratelimit.go new file mode 100644 index 0000000..0c9df34 --- /dev/null +++ b/init/ratelimit/ratelimit.go @@ -0,0 +1,41 @@ +package ratelimit + +import ( + "context" + "net/http" + "time" + + "github.com/gin-gonic/gin" + "go.uber.org/ratelimit" +) + +var rl = ratelimit.New(20, ratelimit.WithoutSlack) + +func RateLimitMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + // 1. 给请求设置超时(如3秒,超过3秒未拿到令牌直接返回) + ctx, cancel := context.WithTimeout(c.Request.Context(), 3*time.Second) + defer cancel() + + // 2. 用带超时的方式获取令牌(需自定义限流器的 Take 逻辑) + ch := make(chan struct{}, 1) + go func() { + rl.Take() // 阻塞拿令牌 + ch <- struct{}{} // 拿到令牌后通知主协程 + }() + + select { + case <-ch: + // 3. 拿到令牌,继续执行 + c.Next() + case <-ctx.Done(): + // 4. 超时未拿到令牌,返回请求频繁 + c.JSON(http.StatusTooManyRequests, gin.H{ + "code": 429, + "message": "请求过于频繁,请稍后重试", + "reason": "等待令牌超时", + }) + c.Abort() + } + } +}