diff --git a/middleware/model-rate-limit.go b/middleware/model-rate-limit.go index 78d8f619d2..22d92e226f 100644 --- a/middleware/model-rate-limit.go +++ b/middleware/model-rate-limit.go @@ -79,14 +79,13 @@ func recordRedisRequest(ctx context.Context, rdb *redis.Client, key string, maxC } // Redis限流处理器 -func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { +func redisRateLimitHandler(subject string, duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { return func(c *gin.Context) { - userId := strconv.Itoa(c.GetInt("id")) ctx := context.Background() rdb := common.RDB // 1. 检查成功请求数限制 - successKey := fmt.Sprintf("rateLimit:%s:%s", ModelRequestRateLimitSuccessCountMark, userId) + successKey := fmt.Sprintf("rateLimit:%s:%s", ModelRequestRateLimitSuccessCountMark, subject) allowed, err := checkRedisRateLimit(ctx, rdb, successKey, successMaxCount, duration) if err != nil { fmt.Println(common.LogText("failed to check the successful request limit: %s", err.Error())) @@ -100,7 +99,7 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g //2.检查总请求数限制并记录总请求（当totalMaxCount为0时会自动跳过，使用令牌桶限流器 if totalMaxCount > 0 { - totalKey := fmt.Sprintf("rateLimit:%s", userId) + totalKey := fmt.Sprintf("rateLimit:%s", subject) // 初始化 tb := limiter.New(ctx, rdb) allowed, err = tb.Allow( @@ -134,13 +133,12 @@ func redisRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) g } // 内存限流处理器 -func memoryRateLimitHandler(duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { +func memoryRateLimitHandler(subject string, duration int64, totalMaxCount, successMaxCount int) gin.HandlerFunc { inMemoryRateLimiter.Init(time.Duration(setting.ModelRequestRateLimitDurationMinutes) * time.Minute) return func(c *gin.Context) { - userId := strconv.Itoa(c.GetInt("id")) - totalKey := ModelRequestRateLimitCountMark + userId - successKey := ModelRequestRateLimitSuccessCountMark + userId + totalKey := ModelRequestRateLimitCountMark + subject + successKey := ModelRequestRateLimitSuccessCountMark + subject // 1. 检查总请求数限制（当totalMaxCount为0时跳过） if totalMaxCount > 0 && !inMemoryRateLimiter.Request(totalKey, totalMaxCount, duration) { @@ -191,18 +189,21 @@ func ModelRequestRateLimit() func(c *gin.Context) { group = common.GetContextKeyString(c, constant.ContextKeyUserGroup) } - //获取分组的限流配置 + // 获取分组的限流配置。配置了独立限流的分组按用户和分组单独计数， + // 其余分组共用该用户的全局计数 + subject := strconv.Itoa(c.GetInt("id")) groupTotalCount, groupSuccessCount, found := setting.GetGroupRateLimit(group) if found { totalMaxCount = groupTotalCount successMaxCount = groupSuccessCount + subject += ":" + group } // 根据存储类型选择并执行限流处理器 if common.RedisEnabled { - redisRateLimitHandler(duration, totalMaxCount, successMaxCount)(c) + redisRateLimitHandler(subject, duration, totalMaxCount, successMaxCount)(c) } else { - memoryRateLimitHandler(duration, totalMaxCount, successMaxCount)(c) + memoryRateLimitHandler(subject, duration, totalMaxCount, successMaxCount)(c) } } } diff --git a/middleware/model_rate_limit_test.go b/middleware/model_rate_limit_test.go index 45cc9e2d19..38bae0de7b 100644 --- a/middleware/model_rate_limit_test.go +++ b/middleware/model_rate_limit_test.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "strconv" "sync/atomic" "testing" "time" @@ -49,14 +50,14 @@ func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) { for _, backend := range []string{"memory", "redis"} { for _, totalLimit := range []int{0, 2} { t.Run(fmt.Sprintf("%s/total=%d", backend, totalLimit), func(t *testing.T) { - userID := 7200000 + int(modelRateLimitTestUsers.Add(1)) - handler := memoryRateLimitHandler(60, totalLimit, 1) + subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1))) + handler := memoryRateLimitHandler(subject, 60, totalLimit, 1) if backend == "redis" { useRateLimitMiniRedis(t) - handler = redisRateLimitHandler(60, totalLimit, 1) + handler = redisRateLimitHandler(subject, 60, totalLimit, 1) } router := gin.New() - router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, handler, func(c *gin.Context) { + router.GET("/:outcome", handler, func(c *gin.Context) { status := relaycommon.NewStreamStatus() if c.Param("outcome") == "failed" { status.MarkFailed("server_error", "", 0) @@ -80,10 +81,10 @@ func TestModelRateLimitStreamFailuresDoNotConsumeSuccessLimit(t *testing.T) { func TestModelMemoryRateLimitReservesConcurrentSuccessAdmission(t *testing.T) { require.NoError(t, i18n.Init()) - userID := 7200000 + int(modelRateLimitTestUsers.Add(1)) + subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1))) entered, release, finished := make(chan struct{}), make(chan struct{}), make(chan struct{}) router := gin.New() - router.GET("/:outcome", func(c *gin.Context) { c.Set("id", userID) }, memoryRateLimitHandler(60, 0, 1), func(c *gin.Context) { + router.GET("/:outcome", memoryRateLimitHandler(subject, 60, 0, 1), func(c *gin.Context) { if c.Param("outcome") == "failed" { close(entered) <-release @@ -118,12 +119,11 @@ func TestModelMemoryRateLimitRejectionsReturnOpenAIError(t *testing.T) { {name: "success", successLimit: 1, messageKey: i18n.MsgRateLimitReached}, } { t.Run(tc.name, func(t *testing.T) { - userID := 7200000 + int(modelRateLimitTestUsers.Add(1)) + subject := strconv.Itoa(7200000 + int(modelRateLimitTestUsers.Add(1))) router := gin.New() router.GET("/", func(c *gin.Context) { - c.Set("id", userID) c.Set(common.RequestIdKey, "req-limited") - }, memoryRateLimitHandler(60, tc.totalLimit, tc.successLimit), func(c *gin.Context) { + }, memoryRateLimitHandler(subject, 60, tc.totalLimit, tc.successLimit), func(c *gin.Context) { c.Status(http.StatusOK) }) require.Equal(t, http.StatusOK, performRateLimitRequest(router, "/", "127.0.0.1:1000").Code) @@ -141,3 +141,59 @@ func TestModelMemoryRateLimitRejectionsReturnOpenAIError(t *testing.T) { }) } } + +// A group with its own limit counts each user's requests on its own; groups +// without one keep sharing the user's global limit. +func TestModelRateLimitCountsConfiguredGroupsSeparately(t *testing.T) { + require.NoError(t, i18n.Init()) + previousEnabled, previousMinutes := setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes + previousTotal, previousSuccess := setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount + previousGroups, previousRedisEnabled := setting.ModelRequestRateLimitGroup, common.RedisEnabled + t.Cleanup(func() { + setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = previousEnabled, previousMinutes + setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = previousTotal, previousSuccess + setting.ModelRequestRateLimitGroup, common.RedisEnabled = previousGroups, previousRedisEnabled + }) + setting.ModelRequestRateLimitEnabled, setting.ModelRequestRateLimitDurationMinutes = true, 1 + common.RedisEnabled = false + + // Groups a and b have their own limits; c and d have none and share the global one. + steps := []struct { + group string + want int + }{ + {"a", http.StatusOK}, {"a", http.StatusTooManyRequests}, + {"b", http.StatusOK}, {"b", http.StatusTooManyRequests}, + {"c", http.StatusOK}, {"d", http.StatusTooManyRequests}, + } + // The Redis token bucket binds to the first Redis client it sees and another + // test already drives it, so the total limit runs on the memory backend only. + for _, tc := range []struct { + backend, limit string + counts [2]int // [total, success] for groups a and b and for the global limit + }{ + {"memory", "success", [2]int{0, 1}}, + {"redis", "success", [2]int{0, 1}}, + {"memory", "total", [2]int{1, 1000}}, + } { + t.Run(tc.backend+"/"+tc.limit, func(t *testing.T) { + if tc.backend == "redis" { + useRateLimitMiniRedis(t) + } + setting.ModelRequestRateLimitCount, setting.ModelRequestRateLimitSuccessCount = tc.counts[0], tc.counts[1] + setting.ModelRequestRateLimitGroup = map[string][2]int{"a": tc.counts, "b": tc.counts} + userID := 7200000 + int(modelRateLimitTestUsers.Add(1)) + router := gin.New() + router.GET("/:group", func(c *gin.Context) { + c.Set("id", userID) + common.SetContextKey(c, constant.ContextKeyTokenGroup, c.Param("group")) + }, ModelRequestRateLimit(), func(c *gin.Context) { + c.Status(http.StatusOK) + }) + + for _, step := range steps { + assert.Equal(t, step.want, performRateLimitRequest(router, "/"+step.group, "127.0.0.1:1000").Code, "group %s", step.group) + } + }) + } +}