package middleware import ( "context" "fmt" "net/http" "strconv" "sync" "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/common/limiter" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" ) const ( UserModelRateLimitMark = "UMRL" ) var userModelRateLimiter common.InMemoryRateLimiter // userModelRpmCache 非 Redis 模式下的进程内缓存,避免每次请求打 DB var userModelRpmCache sync.Map // key: "userId:model" -> rpmCacheEntry type rpmCacheEntry struct { rpm int expiresAt int64 } const rpmCacheTTLSeconds = 300 // 5 分钟,与 Redis 路径一致 // loadUserModelRpm 查询指定用户+模型的 RPM 限制。 // Redis 可用时使用 Hash 缓存批量回填;Redis 不可用时直接查 DB。 // 返回 0 表示不限制(无配置)。 func loadUserModelRpm(userId int, modelName string) int { if common.RedisEnabled { return loadUserModelRpmFromRedis(userId, modelName) } return loadUserModelRpmFromMemory(userId, modelName) } func loadUserModelRpmFromMemory(userId int, modelName string) int { cacheKey := fmt.Sprintf("%d:%s", userId, modelName) now := time.Now().Unix() if v, ok := userModelRpmCache.Load(cacheKey); ok { entry := v.(rpmCacheEntry) if now < entry.expiresAt { return entry.rpm } userModelRpmCache.Delete(cacheKey) } rpm, _ := model.GetUserModelRpm(userId, modelName) userModelRpmCache.Store(cacheKey, rpmCacheEntry{rpm: rpm, expiresAt: now + rpmCacheTTLSeconds}) return rpm } func loadUserModelRpmFromRedis(userId int, modelName string) int { ctx := context.Background() rdb := common.RDB hashKey := fmt.Sprintf("user_model_rate_limit:%d", userId) // 单次 HMGet 同时查目标模型和哨兵字段 vals, err := rdb.HMGet(ctx, hashKey, modelName, "__loaded").Result() if err == nil && len(vals) == 2 { // vals[0] = 模型 RPM,vals[1] = 哨兵 if vals[0] != nil { if s, ok := vals[0].(string); ok { if rpm, parseErr := strconv.Atoi(s); parseErr == nil { return rpm } } } // 哨兵存在说明已加载过,该模型无配置 if vals[1] != nil { return 0 } } // 首次查询:从 DB 加载并批量回填(含哨兵字段,防止未配置用户每次打 DB) limits, dbErr := model.GetUserModelRateLimits(userId) if dbErr != nil { return 0 } pipe := rdb.Pipeline() for _, l := range limits { pipe.HSet(ctx, hashKey, l.Model, l.Rpm) } pipe.HSet(ctx, hashKey, "__loaded", "1") pipe.Expire(ctx, hashKey, 5*time.Minute) if _, err := pipe.Exec(ctx); err != nil { common.SysError(fmt.Sprintf("user model rate limit: failed to cache rpm for user %d: %v", userId, err)) rdb.Del(ctx, hashKey) } for _, l := range limits { if l.Model == modelName { return l.Rpm } } return 0 } // UserModelRateLimit 用户-模型维度的 RPM 速率限制中间件。 // 依赖 context 中的 "id"(userId int)和 "original_model"(modelName string)。 func UserModelRateLimit() gin.HandlerFunc { duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute userModelRateLimiter.Init(duration) return func(c *gin.Context) { // 1. 全局开关 if !setting.ModelRequestRateLimitEnabled { c.Next() return } // 2. 读取 userId / modelName userId := c.GetInt("id") modelName := c.GetString("original_model") if userId == 0 || modelName == "" { c.Next() return } // 3. 查 RPM 配置 rpm := loadUserModelRpm(userId, modelName) if rpm == 0 { c.Next() return } duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute if common.RedisEnabled { userModelRedisRateLimit(c, userId, modelName, rpm, duration) } else { userModelMemoryRateLimit(c, userId, modelName, rpm, duration) } } } // userModelRedisRateLimit Redis 令牌桶限流 func userModelRedisRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) { ctx := context.Background() rdb := common.RDB key := fmt.Sprintf("rateLimit:%s:%d:%s", UserModelRateLimitMark, userId, modelName) durationSeconds := int64(duration.Seconds()) tb := limiter.New(ctx, rdb) allowed, err := tb.Allow(ctx, key, limiter.WithCapacity(int64(rpm)*durationSeconds), limiter.WithRate(int64(rpm)), limiter.WithRequested(durationSeconds), ) if err != nil { // Redis 异常时降级到内存限流 userModelMemoryRateLimit(c, userId, modelName, rpm, duration) return } if !allowed { abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)", setting.ModelRequestRateLimitDurationMinutes, rpm, modelName)) return } c.Next() // 请求失败时退还令牌 if c.Writer.Status() >= 400 { _ = tb.Refund(ctx, key, durationSeconds, int64(rpm)*durationSeconds) } } // userModelMemoryRateLimit 内存滑动窗口限流(参考 memoryRateLimitHandler) func userModelMemoryRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) { key := fmt.Sprintf("%s%d:%s", UserModelRateLimitMark, userId, modelName) durationSeconds := int64(duration.Seconds()) if !userModelRateLimiter.Request(key, rpm, durationSeconds) { abortWithOpenAiMessage(c, http.StatusTooManyRequests, fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)", setting.ModelRequestRateLimitDurationMinutes, rpm, modelName)) return } c.Next() // 请求失败时退还配额 if c.Writer.Status() >= 400 { userModelRateLimiter.Refund(key) } }