25w30e

rate/rate.go

@@ -1,13 +1,14 @@
 package rate
 
 import (
+	"sync"
 	"time"
 
 	"github.com/WJQSERVER-STUDIO/go-utils/logger"
 	"golang.org/x/time/rate"
 )
 
-// 日志输出
+// 日志模块
 var (
 	logw       = logger.Logw
 	logDump    = logger.LogDump
@@ -17,49 +18,90 @@ var (
 	logError   = logger.LogError
 )
 
-// 总体限流器
+// RateLimiter 总体限流器
 type RateLimiter struct {
 	limiter *rate.Limiter
 }
 
-// 基于IP的限流器
-type IPRateLimiter struct {
-	limiters map[string]*RateLimiter
-	limit    int
-	burst    int
-	duration time.Duration
-}
-
+// New 创建一个总体限流器
 func New(limit int, burst int, duration time.Duration) *RateLimiter {
+	if limit <= 0 {
+		limit = 1
+		logWarning("rate limit per minute must be positive, setting to 1")
+	}
+	if burst <= 0 {
+		burst = 1
+		logWarning("rate limit burst must be positive, setting to 1")
+	}
+
+	rateLimit := rate.Limit(float64(limit) / duration.Seconds())
+
 	return &RateLimiter{
-		limiter: rate.NewLimiter(rate.Limit(float64(limit)/duration.Seconds()), burst),
+		limiter: rate.NewLimiter(rateLimit, burst),
 	}
 }
 
+// Allow 检查是否允许请求通过
 func (rl *RateLimiter) Allow() bool {
 	return rl.limiter.Allow()
 }
 
-func NewIPRateLimiter(limit int, burst int, duration time.Duration) *IPRateLimiter {
+// IPRateLimiter 基于IP的限流器
+type IPRateLimiter struct {
+	limiters map[string]*RateLimiter // 用户级限流器 map
+	mu       sync.RWMutex            // 保护 limiters map
+	limit    int                     // 每 duration 时间段内允许的请求数
+	burst    int                     // 突发请求数
+	duration time.Duration           // 限流周期
+}
+
+// NewIPRateLimiter 创建一个基于IP的限流器
+func NewIPRateLimiter(ipLimit int, ipBurst int, duration time.Duration) *IPRateLimiter {
+	if ipLimit <= 0 {
+		ipLimit = 1
+		logWarning("IP rate limit per minute must be positive, setting to 1")
+	}
+	if ipBurst <= 0 {
+		ipBurst = 1
+		logWarning("IP rate limit burst must be positive, setting to 1")
+	}
+
+	logInfo("IP Rate Limiter initialized with limit: %d, burst: %d, duration: %v", ipLimit, ipBurst, duration)
+
 	return &IPRateLimiter{
 		limiters: make(map[string]*RateLimiter),
-		limit:    limit,
-		burst:    burst,
+		limit:    ipLimit,
+		burst:    ipBurst,
 		duration: duration,
 	}
 }
 
+// Allow 检查给定IP的请求是否允许通过
 func (rl *IPRateLimiter) Allow(ip string) bool {
 	if ip == "" {
-		logWarning("empty ip")
+		logWarning("empty ip for rate limiting")
 		return false
 	}
 
-	limiter, ok := rl.limiters[ip]
-	if !ok {
-		// 创建新的 RateLimiter 并存储
-		limiter = New(rl.limit, rl.burst, rl.duration)
-		rl.limiters[ip] = limiter
+	// 使用读锁快速查找
+	rl.mu.RLock()
+	limiter, found := rl.limiters[ip]
+	rl.mu.RUnlock()
+
+	if found {
+		return limiter.Allow()
 	}
+
+	// 未找到，获取写锁来创建和添加
+	rl.mu.Lock()
+	// 双重检查
+	limiter, found = rl.limiters[ip]
+	if !found {
+		newL := New(rl.limit, rl.burst, rl.duration)
+		rl.limiters[ip] = newL
+		limiter = newL
+	}
+	rl.mu.Unlock()
+
 	return limiter.Allow()
 }