|
- 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)
- }
- }
|