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.
2026-05-20 14:40:36 +08:00 · 2026-05-20 14:40:36 +08:00 · 3a25007df2
commit 3a25007df2
parent f069f98b35
7 changed files with 352 additions and 36 deletions
--- 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
 		}
--- a/controller/relay_concurrency_test.go
+++ 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)
+}
--- a/middleware/concurrency.go
+++ 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()
+	}
+}
--- a/middleware/concurrency_test.go
+++ 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)
+}
--- 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 {
--- 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)
--- 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)