添加令牌桶中间件

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()
+		}
+	}
+}