init

internal/api/middleware/ratelimit.go

@@ -0,0 +1,103 @@
+package middleware
+
+import (
+	"net/http"
+	"qd-sc/internal/model"
+	"sync/atomic"
+	"time"
+
+	"github.com/gin-gonic/gin"
+)
+
+// RateLimiter 基于令牌桶的限流器（使用原子操作）
+type RateLimiter struct {
+	tokens     int64 // 当前令牌数
+	capacity   int64 // 桶容量（最大突发请求数）
+	refillRate int64 // 每秒补充的令牌数（持续QPS）
+	lastRefill int64 // 上次补充时间（纳秒时间戳）
+
+	now func() int64 // 便于测试注入（UnixNano）
+}
+
+// NewRateLimiter 创建限流器
+// capacity: 桶容量（最大突发请求数）
+// refillRate: 每秒补充的令牌数（持续QPS）
+func NewRateLimiter(capacity, refillRate int) *RateLimiter {
+	return &RateLimiter{
+		tokens:     int64(capacity),
+		capacity:   int64(capacity),
+		refillRate: int64(refillRate),
+		lastRefill: time.Now().UnixNano(),
+		now: func() int64 {
+			return time.Now().UnixNano()
+		},
+	}
+}
+
+// Allow 尝试消耗一个令牌（使用CAS无锁算法）
+func (rl *RateLimiter) Allow() bool {
+	now := rl.now()
+
+	for {
+		// 读取当前状态
+		currentTokens := atomic.LoadInt64(&rl.tokens)
+		lastRefill := atomic.LoadInt64(&rl.lastRefill)
+
+		// 计算应该补充的令牌
+		elapsed := now - lastRefill
+		if elapsed < 0 {
+			// 时钟回拨等极端情况：不补充
+			elapsed = 0
+		}
+
+		// 安全计算：避免 elapsed * refillRate 直接相乘造成溢出
+		// tokensToAdd = floor(elapsed_ns * refillRate_per_sec / 1e9)
+		secPart := elapsed / int64(time.Second)  // elapsed 秒
+		nsecPart := elapsed % int64(time.Second) // 剩余纳秒
+		tokensToAdd := secPart*rl.refillRate + (nsecPart*rl.refillRate)/int64(time.Second)
+
+		newTokens := currentTokens
+		if tokensToAdd > 0 {
+			newTokens = currentTokens + tokensToAdd
+			if newTokens > rl.capacity {
+				newTokens = rl.capacity
+			}
+		}
+
+		// 检查是否有令牌可用
+		if newTokens < 1 {
+			return false
+		}
+
+		// 尝试消耗一个令牌
+		if atomic.CompareAndSwapInt64(&rl.tokens, currentTokens, newTokens-1) {
+			// 更新最后补充时间
+			if tokensToAdd > 0 {
+				atomic.StoreInt64(&rl.lastRefill, now)
+			}
+			return true
+		}
+
+		// CAS失败，重试
+	}
+}
+
+// RateLimit 限流中间件
+func RateLimit(capacity, refillRate int) gin.HandlerFunc {
+	limiter := NewRateLimiter(capacity, refillRate)
+
+	return func(c *gin.Context) {
+		if !limiter.Allow() {
+			c.JSON(http.StatusTooManyRequests, model.ErrorResponse{
+				Error: model.ErrorDetail{
+					Message: "请求过于频繁，请稍后再试",
+					Type:    "rate_limit_exceeded",
+				},
+			})
+			c.Abort()
+			return
+		}
+
+		c.Next()
+	}
+}