Refactor: Optimize the token bucket algorithm, specifically the New method in common/imiterlimiter.go.

Solution: Remove Redis ping. When printing exceptions, use SysLog to print and add additional logging information.
This commit is contained in:
霍雨佳
2025-04-16 16:36:07 +08:00
parent eb75ff232f
commit e385e347ea
2 changed files with 6 additions and 11 deletions

View File

@@ -5,6 +5,7 @@ import (
_ "embed"
"fmt"
"github.com/go-redis/redis/v8"
"one-api/common"
"sync"
)
@@ -23,19 +24,13 @@ var (
func New(ctx context.Context, r *redis.Client) *RedisLimiter {
once.Do(func() {
client := r
_, err := client.Ping(ctx).Result()
if err != nil {
panic(err) // 或者处理连接错误
}
// 预加载脚本
limitSHA, err := client.ScriptLoad(ctx, rateLimitScript).Result()
limitSHA, err := r.ScriptLoad(ctx, rateLimitScript).Result()
if err != nil {
fmt.Println(err)
common.SysLog(fmt.Sprintf("Failed to load rate limit script: %v", err))
}
instance = &RedisLimiter{
client: client,
client: r,
limitScriptSHA: limitSHA,
}
})

View File

@@ -91,9 +91,9 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g
abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("您已达到请求数限制:%d分钟内最多请求%d次", setting.ModelRequestRateLimitDurationMinutes, successMaxCount))
return
}
//检查总请求数限制并记录总请求当totalMaxCount为0时会自动跳过使用令牌桶限流器
//2.检查总请求数限制并记录总请求当totalMaxCount为0时会自动跳过使用令牌桶限流器
totalKey := fmt.Sprintf("rateLimit:%s", userId)
//allowed, err = checkRedisRateLimit(ctx, rdb, totalKey, totalMaxCount, duration)
// 初始化
tb := limiter.New(ctx, rdb)
allowed, err = tb.Allow(