File size: 2,975 Bytes
daa8246 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 | package service
import (
"fmt"
"strconv"
"sync"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/bytedance/gopkg/util/gopool"
)
// notifyLimitStore is used for in-memory rate limiting when Redis is disabled
var (
notifyLimitStore sync.Map
cleanupOnce sync.Once
)
type limitCount struct {
Count int
Timestamp time.Time
}
func getDuration() time.Duration {
minute := constant.NotificationLimitDurationMinute
return time.Duration(minute) * time.Minute
}
// startCleanupTask starts a background task to clean up expired entries
func startCleanupTask() {
gopool.Go(func() {
for {
time.Sleep(time.Hour)
now := time.Now()
notifyLimitStore.Range(func(key, value interface{}) bool {
if limit, ok := value.(limitCount); ok {
if now.Sub(limit.Timestamp) >= getDuration() {
notifyLimitStore.Delete(key)
}
}
return true
})
}
})
}
// CheckNotificationLimit checks if the user has exceeded their notification limit
// Returns true if the user can send notification, false if limit exceeded
func CheckNotificationLimit(userId int, notifyType string) (bool, error) {
if common.RedisEnabled {
return checkRedisLimit(userId, notifyType)
}
return checkMemoryLimit(userId, notifyType)
}
func checkRedisLimit(userId int, notifyType string) (bool, error) {
key := fmt.Sprintf("notify_limit:%d:%s:%s", userId, notifyType, time.Now().Format("2006010215"))
// Get current count
count, err := common.RedisGet(key)
if err != nil && err.Error() != "redis: nil" {
return false, fmt.Errorf("failed to get notification count: %w", err)
}
// If key doesn't exist, initialize it
if count == "" {
err = common.RedisSet(key, "1", getDuration())
return true, err
}
currentCount, _ := strconv.Atoi(count)
limit := constant.NotifyLimitCount
// Check if limit is already reached
if currentCount >= limit {
return false, nil
}
// Only increment if under limit
err = common.RedisIncr(key, 1)
if err != nil {
return false, fmt.Errorf("failed to increment notification count: %w", err)
}
return true, nil
}
func checkMemoryLimit(userId int, notifyType string) (bool, error) {
// Ensure cleanup task is started
cleanupOnce.Do(startCleanupTask)
key := fmt.Sprintf("%d:%s:%s", userId, notifyType, time.Now().Format("2006010215"))
now := time.Now()
// Get current limit count or initialize new one
var currentLimit limitCount
if value, ok := notifyLimitStore.Load(key); ok {
currentLimit = value.(limitCount)
// Check if the entry has expired
if now.Sub(currentLimit.Timestamp) >= getDuration() {
currentLimit = limitCount{Count: 0, Timestamp: now}
}
} else {
currentLimit = limitCount{Count: 0, Timestamp: now}
}
// Increment count
currentLimit.Count++
// Check against limits
limit := constant.NotifyLimitCount
// Store updated count
notifyLimitStore.Store(key, currentLimit)
return currentLimit.Count <= limit, nil
}
|