From 3a25007df26750133b25b3c76038d0d1220c7f60 Mon Sep 17 00:00:00 2001 From: zizi Date: Wed, 20 May 2026 14:40:36 +0800 Subject: [PATCH] feat: enforce relay concurrency limits Add user RPM and concurrency middleware, enforce channel concurrency for relay selection and retry paths, and wire checks into relay routes. --- controller/relay.go | 104 ++++++++++++++------ controller/relay_concurrency_test.go | 77 +++++++++++++++ middleware/concurrency.go | 48 ++++++++++ middleware/concurrency_test.go | 137 +++++++++++++++++++++++++++ middleware/distributor.go | 8 ++ router/relay-router.go | 8 +- router/video-router.go | 6 +- 7 files changed, 352 insertions(+), 36 deletions(-) create mode 100644 controller/relay_concurrency_test.go create mode 100644 middleware/concurrency.go create mode 100644 middleware/concurrency_test.go diff --git a/controller/relay.go b/controller/relay.go index 5e2db44c..f0305ed3 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -32,6 +32,8 @@ import ( "github.com/gorilla/websocket" ) +var acquireChannelConcurrencyForRelay = service.AcquireChannelConcurrency + func relayHandler(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError { var err *types.NewAPIError switch info.RelayMode { @@ -196,29 +198,39 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { break } - addUsedChannel(c, channel.Id) - bodyStorage, bodyErr := common.GetBodyStorage(c) - if bodyErr != nil { - // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) - if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { - newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) - } else { - newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) - } + releaseRetryChannel, concurrencyErr := acquireRetryChannelConcurrency(channel, retryParam.GetRetry()) + if concurrencyErr != nil { + newAPIError = concurrencyErr break } - c.Request.Body = io.NopCloser(bodyStorage) - switch relayFormat { - case types.RelayFormatOpenAIRealtime: - newAPIError = relay.WssHelper(c, relayInfo) - case types.RelayFormatClaude: - newAPIError = relay.ClaudeHelper(c, relayInfo) - case types.RelayFormatGemini: - newAPIError = geminiRelayHandler(c, relayInfo) - default: - newAPIError = relayHandler(c, relayInfo) - } + func() { + defer releaseRetryChannel() + + addUsedChannel(c, channel.Id) + bodyStorage, bodyErr := common.GetBodyStorage(c) + if bodyErr != nil { + // Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path) + if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { + newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry()) + } else { + newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + return + } + c.Request.Body = io.NopCloser(bodyStorage) + + switch relayFormat { + case types.RelayFormatOpenAIRealtime: + newAPIError = relay.WssHelper(c, relayInfo) + case types.RelayFormatClaude: + newAPIError = relay.ClaudeHelper(c, relayInfo) + case types.RelayFormatGemini: + newAPIError = geminiRelayHandler(c, relayInfo) + default: + newAPIError = relayHandler(c, relayInfo) + } + }() if newAPIError == nil { relayInfo.LastError = nil @@ -321,6 +333,24 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service return channel, nil } +func acquireRetryChannelConcurrency(channel *model.Channel, retry int) (func(), *types.NewAPIError) { + release := func() {} + if retry <= 0 || channel == nil || channel.ConcurrencyLimit <= 0 { + return release, nil + } + + allowed, release := acquireChannelConcurrencyForRelay(channel.Id, channel.ConcurrencyLimit) + if !allowed { + return nil, types.NewErrorWithStatusCode( + errors.New("渠道并发请求数超过限制"), + types.ErrorCodeGetChannelFailed, + http.StatusTooManyRequests, + types.ErrOptionWithSkipRetry(), + ) + } + return release, nil +} + func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool { if openaiErr == nil { return false @@ -534,19 +564,33 @@ func RelayTask(c *gin.Context) { } } - addUsedChannel(c, channel.Id) - bodyStorage, bodyErr := common.GetBodyStorage(c) - if bodyErr != nil { - if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { - taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge) - } else { - taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest) + releaseRetryChannel := func() {} + if lockedCh, locked := relayInfo.LockedChannel.(*model.Channel); !locked || lockedCh == nil { + var concurrencyErr *types.NewAPIError + releaseRetryChannel, concurrencyErr = acquireRetryChannelConcurrency(channel, retryParam.GetRetry()) + if concurrencyErr != nil { + taskErr = service.TaskErrorWrapperLocal(concurrencyErr.Err, "channel_concurrency_limit_exceeded", http.StatusTooManyRequests) + break } - break } - c.Request.Body = io.NopCloser(bodyStorage) - result, taskErr = relay.RelayTaskSubmit(c, relayInfo) + func() { + defer releaseRetryChannel() + + addUsedChannel(c, channel.Id) + bodyStorage, bodyErr := common.GetBodyStorage(c) + if bodyErr != nil { + if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) { + taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusRequestEntityTooLarge) + } else { + taskErr = service.TaskErrorWrapperLocal(bodyErr, "read_request_body_failed", http.StatusBadRequest) + } + return + } + c.Request.Body = io.NopCloser(bodyStorage) + + result, taskErr = relay.RelayTaskSubmit(c, relayInfo) + }() if taskErr == nil { break } diff --git a/controller/relay_concurrency_test.go b/controller/relay_concurrency_test.go new file mode 100644 index 00000000..f9cb5837 --- /dev/null +++ b/controller/relay_concurrency_test.go @@ -0,0 +1,77 @@ +package controller + +import ( + "net/http" + "testing" + + "github.com/QuantumNous/new-api/model" + + "github.com/stretchr/testify/require" +) + +func restoreRelayConcurrencyStub(t *testing.T) { + t.Helper() + + previousAcquire := acquireChannelConcurrencyForRelay + t.Cleanup(func() { + acquireChannelConcurrencyForRelay = previousAcquire + }) +} + +func TestAcquireRetryChannelConcurrencySkipsInitialAttempt(t *testing.T) { + restoreRelayConcurrencyStub(t) + + acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) { + t.Fatalf("expected initial attempt to use distributor-owned acquisition") + return false, func() {} + } + + release, err := acquireRetryChannelConcurrency(&model.Channel{ + Id: 1, + ConcurrencyLimit: 1, + }, 0) + + require.Nil(t, err) + require.NotNil(t, release) + release() +} + +func TestAcquireRetryChannelConcurrencyRejectsExceededRetryChannel(t *testing.T) { + restoreRelayConcurrencyStub(t) + + acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) { + require.Equal(t, 2, channelId) + require.Equal(t, 1, limit) + return false, func() {} + } + + release, err := acquireRetryChannelConcurrency(&model.Channel{ + Id: 2, + ConcurrencyLimit: 1, + }, 1) + + require.Nil(t, release) + require.NotNil(t, err) + require.Equal(t, http.StatusTooManyRequests, err.StatusCode) +} + +func TestAcquireRetryChannelConcurrencyReleasesRetryChannel(t *testing.T) { + restoreRelayConcurrencyStub(t) + + released := false + acquireChannelConcurrencyForRelay = func(channelId, limit int) (bool, func()) { + return true, func() { + released = true + } + } + + release, err := acquireRetryChannelConcurrency(&model.Channel{ + Id: 3, + ConcurrencyLimit: 1, + }, 1) + + require.Nil(t, err) + require.NotNil(t, release) + release() + require.True(t, released) +} diff --git a/middleware/concurrency.go b/middleware/concurrency.go new file mode 100644 index 00000000..e20d9bc0 --- /dev/null +++ b/middleware/concurrency.go @@ -0,0 +1,48 @@ +package middleware + +import ( + "net/http" + + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" +) + +var ( + getUserByIdForConcurrency = model.GetUserById + checkUserRPMForConcurrency = service.CheckUserRPM + acquireUserConcurrencyForCheck = service.AcquireUserConcurrency +) + +func ConcurrencyCheck() gin.HandlerFunc { + return func(c *gin.Context) { + userId := c.GetInt("id") + if userId <= 0 { + c.Next() + return + } + + user, err := getUserByIdForConcurrency(userId, false) + if err != nil || user == nil { + c.Next() + return + } + + if user.RpmLimit > 0 && !checkUserRPMForConcurrency(userId, user.RpmLimit) { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, "请求频率超过限制") + return + } + + if user.ConcurrencyLimit > 0 { + allowed, release := acquireUserConcurrencyForCheck(userId, user.ConcurrencyLimit) + if !allowed { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, "用户并发请求数超过限制") + return + } + defer release() + } + + c.Next() + } +} diff --git a/middleware/concurrency_test.go b/middleware/concurrency_test.go new file mode 100644 index 00000000..e792263b --- /dev/null +++ b/middleware/concurrency_test.go @@ -0,0 +1,137 @@ +package middleware + +import ( + "errors" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/model" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +func restoreConcurrencyStubs(t *testing.T) { + t.Helper() + + previousGetUser := getUserByIdForConcurrency + previousCheckRPM := checkUserRPMForConcurrency + previousAcquire := acquireUserConcurrencyForCheck + t.Cleanup(func() { + getUserByIdForConcurrency = previousGetUser + checkUserRPMForConcurrency = previousCheckRPM + acquireUserConcurrencyForCheck = previousAcquire + }) +} + +func performConcurrencyRequest(handler gin.HandlerFunc, userId int) *httptest.ResponseRecorder { + gin.SetMode(gin.TestMode) + router := gin.New() + router.GET("/test", func(c *gin.Context) { + if userId > 0 { + c.Set("id", userId) + } + }, handler, func(c *gin.Context) { + c.Status(http.StatusNoContent) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodGet, "/test", nil) + router.ServeHTTP(recorder, request) + return recorder +} + +func TestConcurrencyCheckAllowsWhenUserMissingFromContext(t *testing.T) { + restoreConcurrencyStubs(t) + + getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) { + t.Fatalf("expected user lookup to be skipped, got id %d", id) + return nil, nil + } + + recorder := performConcurrencyRequest(ConcurrencyCheck(), 0) + + require.Equal(t, http.StatusNoContent, recorder.Code) +} + +func TestConcurrencyCheckFailOpenWhenUserLookupFails(t *testing.T) { + restoreConcurrencyStubs(t) + + getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) { + return nil, errors.New("database unavailable") + } + + recorder := performConcurrencyRequest(ConcurrencyCheck(), 1) + + require.Equal(t, http.StatusNoContent, recorder.Code) +} + +func TestConcurrencyCheckAllowsWhenLimitsDisabled(t *testing.T) { + restoreConcurrencyStubs(t) + + getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) { + return &model.User{Id: id}, nil + } + checkUserRPMForConcurrency = func(userId, limit int) bool { + t.Fatalf("expected RPM check to be skipped") + return false + } + acquireUserConcurrencyForCheck = func(userId, limit int) (bool, func()) { + t.Fatalf("expected concurrency acquisition to be skipped") + return false, func() {} + } + + recorder := performConcurrencyRequest(ConcurrencyCheck(), 1) + + require.Equal(t, http.StatusNoContent, recorder.Code) +} + +func TestConcurrencyCheckRejectsWhenRPMExceeded(t *testing.T) { + restoreConcurrencyStubs(t) + + getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) { + return &model.User{Id: id, RpmLimit: 1}, nil + } + checkUserRPMForConcurrency = func(userId, limit int) bool { + return false + } + + recorder := performConcurrencyRequest(ConcurrencyCheck(), 1) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) +} + +func TestConcurrencyCheckRejectsWhenUserConcurrencyExceeded(t *testing.T) { + restoreConcurrencyStubs(t) + + getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) { + return &model.User{Id: id, ConcurrencyLimit: 1}, nil + } + acquireUserConcurrencyForCheck = func(userId, limit int) (bool, func()) { + return false, func() {} + } + + recorder := performConcurrencyRequest(ConcurrencyCheck(), 1) + + require.Equal(t, http.StatusTooManyRequests, recorder.Code) +} + +func TestConcurrencyCheckReleasesAfterRequest(t *testing.T) { + restoreConcurrencyStubs(t) + + released := false + getUserByIdForConcurrency = func(id int, selectAll bool) (*model.User, error) { + return &model.User{Id: id, ConcurrencyLimit: 1}, nil + } + acquireUserConcurrencyForCheck = func(userId, limit int) (bool, func()) { + return true, func() { + released = true + } + } + + recorder := performConcurrencyRequest(ConcurrencyCheck(), 1) + + require.Equal(t, http.StatusNoContent, recorder.Code) + require.True(t, released) +} diff --git a/middleware/distributor.go b/middleware/distributor.go index 5635b801..048b3165 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -162,6 +162,14 @@ func Distribute() func(c *gin.Context) { abortWithOpenAiMessage(c, http.StatusTooManyRequests, err.Error()) return } + if channel.ConcurrencyLimit > 0 { + allowed, release := service.AcquireChannelConcurrency(channel.Id, channel.ConcurrencyLimit) + if !allowed { + abortWithOpenAiMessage(c, http.StatusTooManyRequests, "渠道并发请求数超过限制") + return + } + defer release() + } } c.Next() if channel != nil && c.Writer != nil && c.Writer.Status() < http.StatusBadRequest { diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad..7ef9bd54 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -62,7 +62,7 @@ func SetRelayRouter(router *gin.Engine) { playgroundRouter := router.Group("/pg") playgroundRouter.Use(middleware.RouteTag("relay")) playgroundRouter.Use(middleware.SystemPerformanceCheck()) - playgroundRouter.Use(middleware.UserAuth(), middleware.Distribute()) + playgroundRouter.Use(middleware.UserAuth(), middleware.ConcurrencyCheck(), middleware.Distribute()) { playgroundRouter.POST("/chat/completions", controller.Playground) } @@ -70,6 +70,7 @@ func SetRelayRouter(router *gin.Engine) { relayV1Router.Use(middleware.RouteTag("relay")) relayV1Router.Use(middleware.SystemPerformanceCheck()) relayV1Router.Use(middleware.TokenAuth()) + relayV1Router.Use(middleware.ConcurrencyCheck()) relayV1Router.Use(middleware.ModelRequestRateLimit()) { // WebSocket 路由（统一到 Relay） @@ -179,7 +180,7 @@ func SetRelayRouter(router *gin.Engine) { relaySunoRouter := router.Group("/suno") relaySunoRouter.Use(middleware.RouteTag("relay")) relaySunoRouter.Use(middleware.SystemPerformanceCheck()) - relaySunoRouter.Use(middleware.TokenAuth(), middleware.Distribute()) + relaySunoRouter.Use(middleware.TokenAuth(), middleware.ConcurrencyCheck(), middleware.Distribute()) { relaySunoRouter.POST("/submit/:action", controller.RelayTask) relaySunoRouter.POST("/fetch", controller.RelayTaskFetch) @@ -190,6 +191,7 @@ func SetRelayRouter(router *gin.Engine) { relayGeminiRouter.Use(middleware.RouteTag("relay")) relayGeminiRouter.Use(middleware.SystemPerformanceCheck()) relayGeminiRouter.Use(middleware.TokenAuth()) + relayGeminiRouter.Use(middleware.ConcurrencyCheck()) relayGeminiRouter.Use(middleware.ModelRequestRateLimit()) relayGeminiRouter.Use(middleware.Distribute()) { @@ -202,7 +204,7 @@ func SetRelayRouter(router *gin.Engine) { func registerMjRouterGroup(relayMjRouter *gin.RouterGroup) { relayMjRouter.GET("/image/:id", relay.RelayMidjourneyImage) - relayMjRouter.Use(middleware.TokenAuth(), middleware.Distribute()) + relayMjRouter.Use(middleware.TokenAuth(), middleware.ConcurrencyCheck(), middleware.Distribute()) { relayMjRouter.POST("/submit/action", controller.RelayMidjourney) relayMjRouter.POST("/submit/shorten", controller.RelayMidjourney) diff --git a/router/video-router.go b/router/video-router.go index 46145110..14ef4c7e 100644 --- a/router/video-router.go +++ b/router/video-router.go @@ -18,7 +18,7 @@ func SetVideoRouter(router *gin.Engine) { videoV1Router := router.Group("/v1") videoV1Router.Use(middleware.RouteTag("relay")) - videoV1Router.Use(middleware.TokenAuth(), middleware.Distribute()) + videoV1Router.Use(middleware.TokenAuth(), middleware.ConcurrencyCheck(), middleware.Distribute()) { videoV1Router.POST("/video/generations", controller.RelayTask) videoV1Router.GET("/video/generations/:task_id", controller.RelayTaskFetch) @@ -33,7 +33,7 @@ func SetVideoRouter(router *gin.Engine) { klingV1Router := router.Group("/kling/v1") klingV1Router.Use(middleware.RouteTag("relay")) - klingV1Router.Use(middleware.KlingRequestConvert(), middleware.TokenAuth(), middleware.Distribute()) + klingV1Router.Use(middleware.KlingRequestConvert(), middleware.TokenAuth(), middleware.ConcurrencyCheck(), middleware.Distribute()) { klingV1Router.POST("/videos/text2video", controller.RelayTask) klingV1Router.POST("/videos/image2video", controller.RelayTask) @@ -44,7 +44,7 @@ func SetVideoRouter(router *gin.Engine) { // Jimeng official API routes - direct mapping to official API format jimengOfficialGroup := router.Group("jimeng") jimengOfficialGroup.Use(middleware.RouteTag("relay")) - jimengOfficialGroup.Use(middleware.JimengRequestConvert(), middleware.TokenAuth(), middleware.Distribute()) + jimengOfficialGroup.Use(middleware.JimengRequestConvert(), middleware.TokenAuth(), middleware.ConcurrencyCheck(), middleware.Distribute()) { // Maps to: /?Action=CVSync2AsyncSubmitTask&Version=2022-08-31 and /?Action=CVSync2AsyncGetResult&Version=2022-08-31 jimengOfficialGroup.POST("/", controller.RelayTask)