Vous ne pouvez pas sélectionner plus de 25 sujets Les noms de sujets doivent commencer par une lettre ou un nombre, peuvent contenir des tirets ('-') et peuvent comporter jusqu'à 35 caractères.
 
 
 

200 lignes
5.5 KiB

  1. package middleware
  2. import (
  3. "context"
  4. "fmt"
  5. "net/http"
  6. "strconv"
  7. "sync"
  8. "time"
  9. "github.com/QuantumNous/new-api/common"
  10. "github.com/QuantumNous/new-api/common/limiter"
  11. "github.com/QuantumNous/new-api/model"
  12. "github.com/QuantumNous/new-api/setting"
  13. "github.com/gin-gonic/gin"
  14. )
  15. const (
  16. UserModelRateLimitMark = "UMRL"
  17. )
  18. var userModelRateLimiter common.InMemoryRateLimiter
  19. // userModelRpmCache 非 Redis 模式下的进程内缓存,避免每次请求打 DB
  20. var userModelRpmCache sync.Map // key: "userId:model" -> rpmCacheEntry
  21. type rpmCacheEntry struct {
  22. rpm int
  23. expiresAt int64
  24. }
  25. const rpmCacheTTLSeconds = 300 // 5 分钟,与 Redis 路径一致
  26. // loadUserModelRpm 查询指定用户+模型的 RPM 限制。
  27. // Redis 可用时使用 Hash 缓存批量回填;Redis 不可用时直接查 DB。
  28. // 返回 0 表示不限制(无配置)。
  29. func loadUserModelRpm(userId int, modelName string) int {
  30. if common.RedisEnabled {
  31. return loadUserModelRpmFromRedis(userId, modelName)
  32. }
  33. return loadUserModelRpmFromMemory(userId, modelName)
  34. }
  35. func loadUserModelRpmFromMemory(userId int, modelName string) int {
  36. cacheKey := fmt.Sprintf("%d:%s", userId, modelName)
  37. now := time.Now().Unix()
  38. if v, ok := userModelRpmCache.Load(cacheKey); ok {
  39. entry := v.(rpmCacheEntry)
  40. if now < entry.expiresAt {
  41. return entry.rpm
  42. }
  43. userModelRpmCache.Delete(cacheKey)
  44. }
  45. rpm, _ := model.GetUserModelRpm(userId, modelName)
  46. userModelRpmCache.Store(cacheKey, rpmCacheEntry{rpm: rpm, expiresAt: now + rpmCacheTTLSeconds})
  47. return rpm
  48. }
  49. func loadUserModelRpmFromRedis(userId int, modelName string) int {
  50. ctx := context.Background()
  51. rdb := common.RDB
  52. hashKey := fmt.Sprintf("user_model_rate_limit:%d", userId)
  53. // 单次 HMGet 同时查目标模型和哨兵字段
  54. vals, err := rdb.HMGet(ctx, hashKey, modelName, "__loaded").Result()
  55. if err == nil && len(vals) == 2 {
  56. // vals[0] = 模型 RPM,vals[1] = 哨兵
  57. if vals[0] != nil {
  58. if s, ok := vals[0].(string); ok {
  59. if rpm, parseErr := strconv.Atoi(s); parseErr == nil {
  60. return rpm
  61. }
  62. }
  63. }
  64. // 哨兵存在说明已加载过,该模型无配置
  65. if vals[1] != nil {
  66. return 0
  67. }
  68. }
  69. // 首次查询:从 DB 加载并批量回填(含哨兵字段,防止未配置用户每次打 DB)
  70. limits, dbErr := model.GetUserModelRateLimits(userId)
  71. if dbErr != nil {
  72. return 0
  73. }
  74. pipe := rdb.Pipeline()
  75. for _, l := range limits {
  76. pipe.HSet(ctx, hashKey, l.Model, l.Rpm)
  77. }
  78. pipe.HSet(ctx, hashKey, "__loaded", "1")
  79. pipe.Expire(ctx, hashKey, 5*time.Minute)
  80. if _, err := pipe.Exec(ctx); err != nil {
  81. common.SysError(fmt.Sprintf("user model rate limit: failed to cache rpm for user %d: %v", userId, err))
  82. rdb.Del(ctx, hashKey)
  83. }
  84. for _, l := range limits {
  85. if l.Model == modelName {
  86. return l.Rpm
  87. }
  88. }
  89. return 0
  90. }
  91. // UserModelRateLimit 用户-模型维度的 RPM 速率限制中间件。
  92. // 依赖 context 中的 "id"(userId int)和 "original_model"(modelName string)。
  93. func UserModelRateLimit() gin.HandlerFunc {
  94. duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute
  95. userModelRateLimiter.Init(duration)
  96. return func(c *gin.Context) {
  97. // 1. 全局开关
  98. if !setting.ModelRequestRateLimitEnabled {
  99. c.Next()
  100. return
  101. }
  102. // 2. 读取 userId / modelName
  103. userId := c.GetInt("id")
  104. modelName := c.GetString("original_model")
  105. if userId == 0 || modelName == "" {
  106. c.Next()
  107. return
  108. }
  109. // 3. 查 RPM 配置
  110. rpm := loadUserModelRpm(userId, modelName)
  111. if rpm == 0 {
  112. c.Next()
  113. return
  114. }
  115. duration := time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute
  116. if common.RedisEnabled {
  117. userModelRedisRateLimit(c, userId, modelName, rpm, duration)
  118. } else {
  119. userModelMemoryRateLimit(c, userId, modelName, rpm, duration)
  120. }
  121. }
  122. }
  123. // userModelRedisRateLimit Redis 令牌桶限流
  124. func userModelRedisRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) {
  125. ctx := context.Background()
  126. rdb := common.RDB
  127. key := fmt.Sprintf("rateLimit:%s:%d:%s", UserModelRateLimitMark, userId, modelName)
  128. durationSeconds := int64(duration.Seconds())
  129. tb := limiter.New(ctx, rdb)
  130. allowed, err := tb.Allow(ctx, key,
  131. limiter.WithCapacity(int64(rpm)*durationSeconds),
  132. limiter.WithRate(int64(rpm)),
  133. limiter.WithRequested(durationSeconds),
  134. )
  135. if err != nil {
  136. // Redis 异常时降级到内存限流
  137. userModelMemoryRateLimit(c, userId, modelName, rpm, duration)
  138. return
  139. }
  140. if !allowed {
  141. abortWithOpenAiMessage(c, http.StatusTooManyRequests,
  142. fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)",
  143. setting.ModelRequestRateLimitDurationMinutes, rpm, modelName))
  144. return
  145. }
  146. c.Next()
  147. // 请求失败时退还令牌
  148. if c.Writer.Status() >= 400 {
  149. _ = tb.Refund(ctx, key, durationSeconds, int64(rpm)*durationSeconds)
  150. }
  151. }
  152. // userModelMemoryRateLimit 内存滑动窗口限流(参考 memoryRateLimitHandler)
  153. func userModelMemoryRateLimit(c *gin.Context, userId int, modelName string, rpm int, duration time.Duration) {
  154. key := fmt.Sprintf("%s%d:%s", UserModelRateLimitMark, userId, modelName)
  155. durationSeconds := int64(duration.Seconds())
  156. if !userModelRateLimiter.Request(key, rpm, durationSeconds) {
  157. abortWithOpenAiMessage(c, http.StatusTooManyRequests,
  158. fmt.Sprintf("用户模型速率限制:每 %d 分钟最多 %d 次请求 (%s)",
  159. setting.ModelRequestRateLimitDurationMinutes, rpm, modelName))
  160. return
  161. }
  162. c.Next()
  163. // 请求失败时退还配额
  164. if c.Writer.Status() >= 400 {
  165. userModelRateLimiter.Refund(key)
  166. }
  167. }