oadmin/new-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Update rate limit configuration to use dynamic expiration duration

Browse Source
This commit is contained in:
1808837298@qq.com
2025-03-02 17:34:39 +08:00
parent 2ca0d7246d
commit 4084b18071
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
middleware/model-rate-limit.go
View File

@@ -51,7 +51,7 @@ func checkRedisRateLimit(ctx context.Context, rdb *redis.Client, key string, max
// 如果在时间窗口内已达到限制，拒绝请求 // 如果在时间窗口内已达到限制，拒绝请求
subTime := nowTime.Sub(oldTime).Seconds() subTime := nowTime.Sub(oldTime).Seconds()
if int64(subTime) < duration { if int64(subTime) < duration {
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) rdb.Expire(ctx, key, time.Duration(setting.ModelRequestRateLimitDurationMinutes)*time.Minute)
return false, nil return false, nil
} }
@@ -68,7 +68,7 @@ func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxC
now := time.Now().Format(timeFormat) now := time.Now().Format(timeFormat)
rdb.LPush(ctx, key, now) rdb.LPush(ctx, key, now)
rdb.LTrim(ctx, key, 0, int64(maxCount-1)) rdb.LTrim(ctx, key, 0, int64(maxCount-1))
rdb.Expire(ctx, key, common.RateLimitKeyExpirationDuration) rdb.Expire(ctx, key, time.Duration(setting.ModelRequestRateLimitDurationMinutes)*time.Minute)
} }
// Redis限流处理器 // Redis限流处理器
@@ -118,7 +118,7 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g
// 内存限流处理器 // 内存限流处理器
func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc {
inMemoryRateLimiter.Init(common.RateLimitKeyExpirationDuration) inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute)
return func(c *gin.Context) { return func(c *gin.Context) {
userId := strconv.Itoa(c.GetInt("id")) userId := strconv.Itoa(c.GetInt("id"))
@@ -161,6 +161,7 @@ func ModelRequestRateLimit() func(c *gin.Context) {
// 计算限流参数 // 计算限流参数
duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60) duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60)
totalMaxCount := setting.ModelRequestRateLimitCount totalMaxCount := setting.ModelRequestRateLimitCount
successMaxCount := setting.ModelRequestRateLimitSuccessCount successMaxCount := setting.ModelRequestRateLimitSuccessCount
// 根据存储类型选择限流处理器 // 根据存储类型选择限流处理器