refactor: Improve model request rate limit middleware execution

This commit is contained in:
1808837298@qq.com
2025-03-06 16:32:11 +08:00
parent e3f9ef1894
commit b841ce006f

View File

@@ -153,21 +153,23 @@ func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int)
// ModelRequestRateLimit 模型请求限流中间件 // ModelRequestRateLimit 模型请求限流中间件
func ModelRequestRateLimit() func(c *gin.Context) { func ModelRequestRateLimit() func(c *gin.Context) {
// 如果未启用限流,直接放行 return func(c *gin.Context) {
if !setting.ModelRequestRateLimitEnabled { // 在每个请求时检查是否启用限流
return defNext if !setting.ModelRequestRateLimitEnabled {
} c.Next()
return
}
// 计算限流参数 // 计算限流参数
duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60) duration := int64(setting.ModelRequestRateLimitDurationMinutes * 60)
totalMaxCount := setting.ModelRequestRateLimitCount totalMaxCount := setting.ModelRequestRateLimitCount
successMaxCount := setting.ModelRequestRateLimitSuccessCount
successMaxCount := setting.ModelRequestRateLimitSuccessCount // 根据存储类型选择并执行限流处理器
if common.RedisEnabled {
// 根据存储类型选择限流处理器 redisRateLimitHandler(duration, totalMaxCount, successMaxCount)(c)
if common.RedisEnabled { } else {
return redisRateLimitHandler(duration, totalMaxCount, successMaxCount) memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)(c)
} else { }
return memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)
} }
} }