opensource/sub2api
3
0
Fork 0
mirror of https://github.com/Wei-Shaw/sub2api.git synced 2026-08-31 01:13:06 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: move user wait queue accounting off hot path

Browse Source
This commit is contained in:
jjaw
2026-06-14 02:40:23 +08:00
committed by shaw
parent 87d0f8da66
commit b0579c4891
8 changed files with 145 additions and 133 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/internal/handler/concurrency_error_response.go
+6
View File
@@ -10,6 +10,12 @@ import (
const statusClientClosedRequest = 499
func concurrencyErrorResponse(err error, slotType string) (int, string, string) {
var waitQueueFullErr *WaitQueueFullError
if errors.As(err, &waitQueueFullErr) {
return http.StatusTooManyRequests, "rate_limit_error",
"Too many pending requests, please retry later"
}
var concurrencyErr *ConcurrencyError
if errors.As(err, &concurrencyErr) {
if concurrencyErr.SlotType != "" {
backend/internal/handler/gateway_handler.go
-27
View File
@@ -211,28 +211,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
// 获取订阅信息（可能为nil）- 提前获取用于后续检查
subscription, _ := middleware2.GetSubscriptionFromContext(c)
// 0. 检查wait队列是否已满
maxWait := service.CalculateMaxWait(subject.Concurrency)
canWait, err := h.concurrencyHelper.IncrementWaitCount(c.Request.Context(), subject.UserID, maxWait)
waitCounted := false
if err != nil {
reqLog.Warn("gateway.user_wait_counter_increment_failed", zap.Error(err))
// On error, allow request to proceed
} else if !canWait {
reqLog.Info("gateway.user_wait_queue_full", zap.Int("max_wait", maxWait))
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
return
}
if err == nil && canWait {
waitCounted = true
}
// Ensure we decrement if we exit before acquiring the user slot.
defer func() {
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
}
}()
// 1. 首先获取用户并发槽位
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
if err != nil {
@@ -240,11 +218,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h.handleConcurrencyError(c, err, "user", streamStarted)
return
}
// User slot acquired: no longer waiting in the queue.
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
waitCounted = false
}
// 在请求结束或 Context 取消时确保释放槽位，避免客户端断开造成泄漏
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
if userReleaseFunc != nil {
backend/internal/handler/gateway_handler_chat_completions.go
-23
View File
@@ -109,35 +109,12 @@ func (h *GatewayHandler) ChatCompletions(c *gin.Context) {
service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds())
// 1. Acquire user concurrency slot
maxWait := service.CalculateMaxWait(subject.Concurrency)
canWait, err := h.concurrencyHelper.IncrementWaitCount(c.Request.Context(), subject.UserID, maxWait)
waitCounted := false
if err != nil {
reqLog.Warn("gateway.cc.user_wait_counter_increment_failed", zap.Error(err))
} else if !canWait {
h.chatCompletionsErrorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
return
}
if err == nil && canWait {
waitCounted = true
}
defer func() {
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
}
}()
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
if err != nil {
reqLog.Warn("gateway.cc.user_slot_acquire_failed", zap.Error(err))
h.handleConcurrencyError(c, err, "user", streamStarted)
return
}
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
waitCounted = false
}
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
if userReleaseFunc != nil {
defer userReleaseFunc()
backend/internal/handler/gateway_handler_responses.go
-23
View File
@@ -118,35 +118,12 @@ func (h *GatewayHandler) Responses(c *gin.Context) {
service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds())
// 1. Acquire user concurrency slot
maxWait := service.CalculateMaxWait(subject.Concurrency)
canWait, err := h.concurrencyHelper.IncrementWaitCount(c.Request.Context(), subject.UserID, maxWait)
waitCounted := false
if err != nil {
reqLog.Warn("gateway.responses.user_wait_counter_increment_failed", zap.Error(err))
} else if !canWait {
h.responsesErrorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
return
}
if err == nil && canWait {
waitCounted = true
}
defer func() {
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
}
}()
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, subject.UserID, subject.Concurrency, reqStream, &streamStarted)
if err != nil {
reqLog.Warn("gateway.responses.user_slot_acquire_failed", zap.Error(err))
h.handleConcurrencyError(c, err, "user", streamStarted)
return
}
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(c.Request.Context(), subject.UserID)
waitCounted = false
}
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
if userReleaseFunc != nil {
defer userReleaseFunc()
backend/internal/handler/gateway_helper.go
+26 -1
View File
@@ -127,6 +127,14 @@ func (e *ConcurrencyError) Error() string {
return fmt.Sprintf("%s concurrency limit reached", e.SlotType)
}
type WaitQueueFullError struct {
SlotType string
}
func (e *WaitQueueFullError) Error() string {
return "Too many pending requests, please retry later"
}
// ConcurrencyHelper provides common concurrency slot management for gateway handlers
type ConcurrencyHelper struct {
concurrencyService *service.ConcurrencyService
@@ -220,6 +228,10 @@ func (h *ConcurrencyHelper) TryAcquireAccountSlot(ctx context.Context, accountID
// For streaming requests, sends ping events during the wait.
// streamStarted is updated if streaming response has begun.
func (h *ConcurrencyHelper) AcquireUserSlotWithWait(c *gin.Context, userID int64, maxConcurrency int, isStream bool, streamStarted *bool) (func(), error) {
return h.acquireUserSlotWithWaitTimeout(c, userID, maxConcurrency, maxConcurrencyWait, isStream, streamStarted)
}
func (h *ConcurrencyHelper) acquireUserSlotWithWaitTimeout(c *gin.Context, userID int64, maxConcurrency int, timeout time.Duration, isStream bool, streamStarted *bool) (func(), error) {
ctx := c.Request.Context()
// Try to acquire immediately
@@ -232,8 +244,21 @@ func (h *ConcurrencyHelper) AcquireUserSlotWithWait(c *gin.Context, userID int64
return releaseFunc, nil
}
queueLimit := service.CalculateMaxWait(maxConcurrency) - maxConcurrency
if queueLimit < 1 {
queueLimit = 1
}
canWait, err := h.IncrementWaitCount(ctx, userID, queueLimit)
if err != nil {
return nil, err
}
if !canWait {
return nil, &WaitQueueFullError{SlotType: "user"}
}
defer h.DecrementWaitCount(ctx, userID)
// Need to wait - handle streaming ping if needed
return h.waitForSlotWithPing(c, "user", userID, maxConcurrency, isStream, streamStarted)
return h.waitForSlotWithPingTimeout(c, "user", userID, maxConcurrency, timeout, isStream, streamStarted, false)
}
// AcquireAccountSlotWithWait acquires an account concurrency slot, waiting if necessary.
backend/internal/handler/gateway_helper_hotpath_test.go
+112
View File
@@ -24,6 +24,11 @@ type helperConcurrencyCacheStub struct {
userAcquireCalls int
accountReleaseCalls int
userReleaseCalls int
waitAllowed bool
waitIncrementCalls int
waitDecrementCalls int
waitMaxWait int
waitIncrementHook func()
}
func (s *helperConcurrencyCacheStub) AcquireAccountSlot(ctx context.Context, accountID int64, maxConcurrency int, requestID string) (bool, error) {
@@ -93,10 +98,25 @@ func (s *helperConcurrencyCacheStub) GetUserConcurrency(ctx context.Context, use
}
func (s *helperConcurrencyCacheStub) IncrementWaitCount(ctx context.Context, userID int64, maxWait int) (bool, error) {
s.mu.Lock()
s.waitIncrementCalls++
s.waitMaxWait = maxWait
waitAllowed := s.waitAllowed
hook := s.waitIncrementHook
s.mu.Unlock()
if hook != nil {
hook()
}
if !waitAllowed {
return false, nil
}
return true, nil
}
func (s *helperConcurrencyCacheStub) DecrementWaitCount(ctx context.Context, userID int64) error {
s.mu.Lock()
defer s.mu.Unlock()
s.waitDecrementCalls++
return nil
}
@@ -226,6 +246,98 @@ func TestWaitForSlotWithPingTimeout_AccountAndUserAcquire(t *testing.T) {
})
}
func TestAcquireUserSlotWithWait_ImmediateAcquireSkipsWaitQueue(t *testing.T) {
cache := &helperConcurrencyCacheStub{
userSeq: []bool{true},
}
concurrency := service.NewConcurrencyService(cache)
helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond)
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
streamStarted := false
release, err := helper.acquireUserSlotWithWaitTimeout(c, 202, 3, time.Second, false, &streamStarted)
require.NoError(t, err)
require.NotNil(t, release)
release()
require.Equal(t, 1, cache.userAcquireCalls)
require.Equal(t, 0, cache.waitIncrementCalls)
require.Equal(t, 0, cache.waitDecrementCalls)
require.Equal(t, 1, cache.userReleaseCalls)
}
func TestAcquireUserSlotWithWait_WaitSuccessDecrementsBeforeReturn(t *testing.T) {
cache := &helperConcurrencyCacheStub{
userSeq: []bool{false, true},
waitAllowed: true,
}
concurrency := service.NewConcurrencyService(cache)
helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond)
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
streamStarted := false
release, err := helper.acquireUserSlotWithWaitTimeout(c, 202, 3, time.Second, false, &streamStarted)
require.NoError(t, err)
require.NotNil(t, release)
require.Equal(t, 2, cache.userAcquireCalls)
require.Equal(t, 1, cache.waitIncrementCalls)
require.Equal(t, 20, cache.waitMaxWait)
require.Equal(t, 1, cache.waitDecrementCalls)
release()
require.Equal(t, 1, cache.userReleaseCalls)
}
func TestAcquireUserSlotWithWait_TimeoutDecrementsWaitQueue(t *testing.T) {
cache := &helperConcurrencyCacheStub{
userSeq: []bool{false, false, false},
waitAllowed: true,
}
concurrency := service.NewConcurrencyService(cache)
helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond)
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
streamStarted := false
release, err := helper.acquireUserSlotWithWaitTimeout(c, 202, 3, 30*time.Millisecond, false, &streamStarted)
require.Nil(t, release)
var cErr *ConcurrencyError
require.ErrorAs(t, err, &cErr)
require.True(t, cErr.IsTimeout)
require.Equal(t, 1, cache.waitIncrementCalls)
require.Equal(t, 1, cache.waitDecrementCalls)
require.Equal(t, 0, cache.userReleaseCalls)
}
func TestAcquireUserSlotWithWait_RequestCancelDecrementsWaitQueue(t *testing.T) {
cancelled := make(chan struct{})
var cancel context.CancelFunc
cache := &helperConcurrencyCacheStub{
userSeq: []bool{false, false},
waitAllowed: true,
waitIncrementHook: func() {
cancel()
close(cancelled)
},
}
concurrency := service.NewConcurrencyService(cache)
helper := NewConcurrencyHelper(concurrency, SSEPingFormatNone, 5*time.Millisecond)
c, _ := newHelperTestContext(http.MethodPost, "/v1/messages")
reqCtx, cancelFunc := context.WithCancel(c.Request.Context())
cancel = cancelFunc
defer cancel()
c.Request = c.Request.WithContext(reqCtx)
streamStarted := false
release, err := helper.acquireUserSlotWithWaitTimeout(c, 202, 3, time.Second, false, &streamStarted)
<-cancelled
require.Nil(t, release)
require.ErrorIs(t, err, context.Canceled)
require.Equal(t, 1, cache.waitIncrementCalls)
require.Equal(t, 1, cache.waitDecrementCalls)
require.Equal(t, 0, cache.userReleaseCalls)
}
func TestWaitForSlotWithPingTimeout_TimeoutAndStreamPing(t *testing.T) {
cache := &helperConcurrencyCacheStub{
accountSeq: []bool{false, false, false},
backend/internal/handler/gemini_v1beta_handler.go
-24
View File
@@ -205,26 +205,6 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
// For Gemini native API, do not send Claude-style ping frames.
geminiConcurrency := NewConcurrencyHelper(h.concurrencyHelper.concurrencyService, SSEPingFormatNone, 0)
// 0) wait queue check
maxWait := service.CalculateMaxWait(authSubject.Concurrency)
canWait, err := geminiConcurrency.IncrementWaitCount(c.Request.Context(), authSubject.UserID, maxWait)
waitCounted := false
if err != nil {
reqLog.Warn("gemini.user_wait_counter_increment_failed", zap.Error(err))
} else if !canWait {
reqLog.Info("gemini.user_wait_queue_full", zap.Int("max_wait", maxWait))
googleError(c, http.StatusTooManyRequests, "Too many pending requests, please retry later")
return
}
if err == nil && canWait {
waitCounted = true
}
defer func() {
if waitCounted {
geminiConcurrency.DecrementWaitCount(c.Request.Context(), authSubject.UserID)
}
}()
// 1) user concurrency slot
streamStarted := false
if h.errorPassthroughService != nil {
@@ -236,10 +216,6 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
googleError(c, http.StatusTooManyRequests, err.Error())
return
}
if waitCounted {
geminiConcurrency.DecrementWaitCount(c.Request.Context(), authSubject.UserID)
waitCounted = false
}
// 确保请求取消时也会释放槽位，避免长连接被动中断造成泄漏
userReleaseFunc = wrapReleaseOnDone(c.Request.Context(), userReleaseFunc)
if userReleaseFunc != nil {
backend/internal/handler/openai_gateway_handler.go
+1 -35
View File
@@ -1019,46 +1019,12 @@ func (h *OpenAIGatewayHandler) acquireResponsesUserSlot(
reqLog *zap.Logger,
) (func(), bool) {
ctx := c.Request.Context()
userReleaseFunc, userAcquired, err := h.concurrencyHelper.TryAcquireUserSlot(ctx, userID, userConcurrency)
userReleaseFunc, err := h.concurrencyHelper.AcquireUserSlotWithWait(c, userID, userConcurrency, reqStream, streamStarted)
if err != nil {
reqLog.Warn("openai.user_slot_acquire_failed", zap.Error(err))
h.handleConcurrencyError(c, err, "user", *streamStarted)
return nil, false
}
if userAcquired {
return wrapReleaseOnDone(ctx, userReleaseFunc), true
}
maxWait := service.CalculateMaxWait(userConcurrency)
canWait, waitErr := h.concurrencyHelper.IncrementWaitCount(ctx, userID, maxWait)
if waitErr != nil {
reqLog.Warn("openai.user_wait_counter_increment_failed", zap.Error(waitErr))
// 按现有降级语义：等待计数异常时放行后续抢槽流程
} else if !canWait {
reqLog.Info("openai.user_wait_queue_full", zap.Int("max_wait", maxWait))
h.errorResponse(c, http.StatusTooManyRequests, "rate_limit_error", "Too many pending requests, please retry later")
return nil, false
}
waitCounted := waitErr == nil && canWait
defer func() {
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(ctx, userID)
}
}()
userReleaseFunc, err = h.concurrencyHelper.AcquireUserSlotWithWait(c, userID, userConcurrency, reqStream, streamStarted)
if err != nil {
reqLog.Warn("openai.user_slot_acquire_failed_after_wait", zap.Error(err))
h.handleConcurrencyError(c, err, "user", *streamStarted)
return nil, false
}
// 槽位获取成功后，立刻退出等待计数。
if waitCounted {
h.concurrencyHelper.DecrementWaitCount(ctx, userID)
waitCounted = false
}
return wrapReleaseOnDone(ctx, userReleaseFunc), true
}