fix(grok): release abandoned media slots and preserve video ownership

Claim scheduler-acquired slots before media eligibility checks and task-owner validation. Release abandoned attempts immediately and retain once-only cleanup on cancellation and forwarding completion. Select only the resolved video owner, reuse bounded account waiting without sticky escape, and preserve the task binding and TTL. Keep generation failover and owner/billing isolation intact. Regression tests reproduce unreleased account slots on upstream ab99d56e9 and pass with this change. Validation: go test ./..., unit and integration suites, focused race tests, backend build, and golangci-lint (0 issues).
2026-10-31 00:24:38 +00:00 · 2026-09-05 20:25:08 +08:00
parent ab99d56e96
commit b97a798ebc
5 changed files with 609 additions and 22 deletions
@@ -198,30 +198,57 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
 	}
 	routingStart := time.Now()
 	requiredCapability := grokMediaRequiredCapability(endpoint)
+	var accountReleaseFunc func()
+	releaseAccount := func() {
+		if accountReleaseFunc != nil {
+			accountReleaseFunc()
+			accountReleaseFunc = nil
+		}
+	}
+	defer releaseAccount()

 	for {
+		releaseAccount()
 		if failoverClientGone(c) {
 			return
 		}
-		selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
-			requestCtx,
-			apiKey.GroupID,
-			"",
-			sessionHash,
-			routingModel,
-			failedAccountIDs,
-			service.OpenAIUpstreamTransportHTTPSSE,
-			requiredCapability,
-			false,
-			false,
-			false,
-			service.PlatformGrok,
-		)
+		var selection *service.AccountSelectionResult
+		var scheduleDecision service.OpenAIAccountScheduleDecision
+		if boundLookupAccountID > 0 {
+			selection, scheduleDecision, err = h.gatewayService.SelectGrokMediaVideoRequestAccount(
+				requestCtx, apiKey.GroupID, sessionHash, boundLookupAccountID, routingModel,
+			)
+		} else {
+			selection, scheduleDecision, err = h.gatewayService.SelectAccountWithSchedulerForCapability(
+				requestCtx,
+				apiKey.GroupID,
+				"",
+				sessionHash,
+				routingModel,
+				failedAccountIDs,
+				service.OpenAIUpstreamTransportHTTPSSE,
+				requiredCapability,
+				false,
+				false,
+				false,
+				service.PlatformGrok,
+			)
+		}
+		// Own an eagerly acquired slot before any rejection or eligibility probe.
+		// Forwarding takes over the same once-only release after admission.
+		if selection != nil && selection.Acquired {
+			selection.ReleaseFunc = wrapReleaseOnDone(requestCtx, selection.ReleaseFunc)
+			accountReleaseFunc = selection.ReleaseFunc
+		}
 		if err != nil {
 			if failoverClientGone(c) {
 				reqLog.Info("grok_media.account_select_aborted_client_disconnected", zap.Error(err))
 				return
 			}
+			if boundLookupAccountID > 0 && errors.Is(err, service.ErrNoAvailableAccounts) {
+				h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
+				return
+			}
 			reqLog.Warn("grok_media.account_select_failed",
 				zap.Error(err),
 				zap.Int("excluded_account_count", len(failedAccountIDs)),
@@ -260,6 +287,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
 			h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
 			return
 		}
+		if failoverClientGone(c) {
+			return
+		}
 		if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID {
 			reqLog.Warn("grok_media.video_lookup_bound_account_unavailable",
 				zap.Int64("bound_account_id", boundLookupAccountID),
@@ -282,6 +312,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
 		if endpoint.IsGenerationRequest() {
 			eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account)
 			if !eligible {
+				releaseAccount()
 				mediaEligibilityRejected = true
 				failedAccountIDs[account.ID] = struct{}{}
 				reqLog.Warn("grok_media.account_eligibility_rejected",
@@ -298,10 +329,19 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
 				continue
 			}
 		}
+		if failoverClientGone(c) {
+			return
+		}
 		sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account)
 		setOpsSelectedAccount(c, account.ID, account.Platform)

-		accountReleaseFunc, slotResult := h.acquireResponsesAccountSlot(c, apiKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog)
+		admissionSessionHash := sessionHash
+		if boundLookupAccountID > 0 {
+			// Waiting must not replace the video ownership TTL with text sticky TTL.
+			admissionSessionHash = ""
+		}
+		var slotResult openAISlotAcquireResult
+		accountReleaseFunc, slotResult = h.acquireResponsesAccountSlot(c, apiKey.GroupID, admissionSessionHash, selection, false, &streamStarted, reqLog)
 		if slotResult == openAISlotAcquireProfitVetoed {
 			// 媒体路径已显式豁免利润门（suppress 标记），此分支仅防御性兜底，
 			// 同样受否决上限约束。
@@ -319,11 +359,7 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
 		forwardStart := time.Now()
 		writerSizeBeforeForward := c.Writer.Size()
 		result, err := func() (*service.OpenAIForwardResult, error) {
-			defer func() {
-				if accountReleaseFunc != nil {
-					accountReleaseFunc()
-				}
-			}()
+			defer releaseAccount()
 			return h.gatewayService.ForwardGrokMedia(requestCtx, c, account, endpoint, requestID, body, contentType)
 		}()

@@ -0,0 +1,405 @@
+//go:build unit
+
+package handler
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/Wei-Shaw/sub2api/internal/config"
+	middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/Wei-Shaw/sub2api/internal/testutil"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+type grokMediaSlotsCache struct {
+	testutil.StubConcurrencyCache
+	mu                                             sync.Mutex
+	accounts                                       map[string]int64
+	users                                          map[string]int64
+	acquired, released, userAcquired, userReleased int
+	waiting, denied                                int
+	full                                           bool
+	queueFull                                      bool
+	onWait                                         func()
+	duplicateRelease                               bool
+	peak                                           int
+}
+
+func (s *grokMediaSlotsCache) AcquireAccountSlot(_ context.Context, id int64, _ int, request string) (bool, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.full || s.denied > 0 {
+		if s.denied > 0 {
+			s.denied--
+		}
+		return false, nil
+	}
+	s.accounts[request] = id
+	s.acquired++
+	if len(s.accounts) > s.peak {
+		s.peak = len(s.accounts)
+	}
+	return true, nil
+}
+
+func (s *grokMediaSlotsCache) ReleaseAccountSlot(_ context.Context, _ int64, request string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if _, ok := s.accounts[request]; !ok {
+		s.duplicateRelease = true
+	}
+	delete(s.accounts, request)
+	s.released++
+	return nil
+}
+
+func (s *grokMediaSlotsCache) AcquireUserSlot(_ context.Context, id int64, _ int, request string) (bool, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.users[request] = id
+	s.userAcquired++
+	return true, nil
+}
+
+func (s *grokMediaSlotsCache) ReleaseUserSlot(_ context.Context, _ int64, request string) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if _, ok := s.users[request]; !ok {
+		s.duplicateRelease = true
+	}
+	delete(s.users, request)
+	s.userReleased++
+	return nil
+}
+
+func (s *grokMediaSlotsCache) IncrementAccountWaitCount(context.Context, int64, int) (bool, error) {
+	s.mu.Lock()
+	if s.queueFull {
+		s.mu.Unlock()
+		return false, nil
+	}
+	s.waiting++
+	onWait := s.onWait
+	s.mu.Unlock()
+	if onWait != nil {
+		onWait()
+	}
+	return true, nil
+}
+
+func (s *grokMediaSlotsCache) DecrementAccountWaitCount(context.Context, int64) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.waiting--
+	return nil
+}
+
+func (s *grokMediaSlotsCache) assertReleased(t *testing.T) {
+	t.Helper()
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	require.Empty(t, s.accounts)
+	require.Empty(t, s.users)
+	require.Equal(t, s.acquired, s.released)
+	require.Equal(t, s.userAcquired, s.userReleased)
+	require.False(t, s.duplicateRelease)
+	require.Zero(t, s.waiting)
+}
+
+type grokMediaSlotBindings struct {
+	testutil.StubGatewayCache
+	owner  int64
+	writes int
+	key    string
+	billed map[string]bool
+}
+
+func (s *grokMediaSlotBindings) GetSessionAccountID(_ context.Context, groupID int64, key string) (int64, error) {
+	if groupID != 24 || key != s.key {
+		return 0, service.ErrStickySessionNotFound
+	}
+	return s.owner, nil
+}
+func (s *grokMediaSlotBindings) SetSessionAccountID(_ context.Context, _ int64, key string, owner int64, _ time.Duration) error {
+	s.key, s.owner = key, owner
+	s.writes++
+	return nil
+}
+func (s *grokMediaSlotBindings) RefreshSessionTTL(context.Context, int64, string, time.Duration) error {
+	s.writes++
+	return nil
+}
+func (s *grokMediaSlotBindings) DeleteSessionAccountID(context.Context, int64, string) error {
+	s.writes++
+	return nil
+}
+
+func (s *grokMediaSlotBindings) ClaimGrokVideoBilled(_ context.Context, key string, _ time.Duration) (bool, error) {
+	if s.billed == nil {
+		s.billed = make(map[string]bool)
+	}
+	if s.billed[key] {
+		return false, nil
+	}
+	s.billed[key] = true
+	return true, nil
+}
+
+type grokMediaSlotRepo struct {
+	openAIImagesFailoverAccountRepo
+	mismatch bool
+}
+
+func (r grokMediaSlotRepo) GetByID(ctx context.Context, id int64) (*service.Account, error) {
+	if r.mismatch {
+		return r.openAIImagesFailoverAccountRepo.GetByID(ctx, 2)
+	}
+	return r.openAIImagesFailoverAccountRepo.GetByID(ctx, id)
+}
+
+type grokMediaSlotUpstream struct {
+	service.HTTPUpstream
+	call  func(*http.Request, int64) (*http.Response, error)
+	calls int
+}
+
+func (u *grokMediaSlotUpstream) Do(req *http.Request, _ string, id int64, _ int) (*http.Response, error) {
+	u.calls++
+	return u.call(req, id)
+}
+
+type grokMediaSlotProber func(context.Context, int64) (bool, string, error)
+
+func (p grokMediaSlotProber) ProbeMediaEligibility(ctx context.Context, id int64) (bool, string, error) {
+	return p(ctx, id)
+}
+
+func newGrokMediaSlotHandler(t *testing.T, oauth, mismatch bool) (*OpenAIGatewayHandler, *grokMediaSlotsCache, *grokMediaSlotBindings, *grokMediaSlotUpstream) {
+	t.Helper()
+	gin.SetMode(gin.TestMode)
+	accounts := make([]service.Account, 3)
+	for i := range accounts {
+		accounts[i] = service.Account{ID: int64(i + 1), Platform: service.PlatformGrok, Type: service.AccountTypeAPIKey,
+			Status: service.StatusActive, Schedulable: true, Concurrency: 50, Priority: i,
+			GroupIDs: []int64{24}, Credentials: map[string]any{"api_key": "test-key", "access_token": "test-token"}}
+		if oauth {
+			accounts[i].Type = service.AccountTypeOAuth
+			accounts[i].Credentials["refresh_token"] = "test-refresh"
+			accounts[i].Credentials["expires_at"] = time.Now().Add(7 * 24 * time.Hour).UTC().Format(time.RFC3339)
+		}
+	}
+	slots := &grokMediaSlotsCache{accounts: map[string]int64{}, users: map[string]int64{}}
+	concurrency := service.NewConcurrencyService(slots)
+	bindings := &grokMediaSlotBindings{owner: 1}
+	upstream := &grokMediaSlotUpstream{call: func(*http.Request, int64) (*http.Response, error) {
+		return &http.Response{StatusCode: 200, Header: http.Header{"Content-Type": []string{"application/json"}},
+			Body: io.NopCloser(strings.NewReader(`{"request_id":"task","status":"pending"}`))}, nil
+	}}
+	cfg := &config.Config{RunMode: config.RunModeSimple}
+	cfg.Gateway.Scheduling.StickySessionWaitTimeout = 20 * time.Millisecond
+	cfg.Gateway.Scheduling.StickySessionMaxWaiting = 3
+	cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
+	repo := grokMediaSlotRepo{openAIImagesFailoverAccountRepo: openAIImagesFailoverAccountRepo{accounts: accounts}, mismatch: mismatch}
+	provider := service.NewGrokTokenProvider(repo, nil)
+	if oauth {
+		_, err := provider.GetAccessToken(context.Background(), &accounts[1])
+		require.NoError(t, err)
+	}
+	gateway := service.NewOpenAIGatewayService(repo, nil, nil, nil, nil, nil, bindings, cfg, nil, concurrency, nil, nil, nil, upstream, nil, nil, provider, nil, nil, nil, nil, nil)
+	groupID := int64(24)
+	require.NoError(t, gateway.BindGrokMediaVideoRequestAccount(context.Background(), &groupID, "task", 10, 20, 1))
+	bindings.writes = 0
+	billing := service.NewBillingCacheService(nil, nil, nil, nil, nil, nil, cfg, nil)
+	t.Cleanup(billing.Stop)
+	handler := NewOpenAIGatewayHandler(gateway, concurrency, billing, service.NewAPIKeyService(nil, nil, nil, nil, nil, nil, cfg), nil, nil, nil, nil, cfg)
+	return handler, slots, bindings, upstream
+}
+
+func grokMediaSlotContext(ctx context.Context, generation bool) (*gin.Context, *httptest.ResponseRecorder) {
+	groupID := int64(24)
+	method, path, body := http.MethodGet, "/v1/videos/task", ""
+	if generation {
+		method, path, body = http.MethodPost, "/v1/videos/generations", `{"model":"grok-imagine-video","prompt":"test","duration":6}`
+	}
+	req := httptest.NewRequest(method, path, strings.NewReader(body)).WithContext(ctx)
+	req.Header.Set("Content-Type", "application/json")
+	if !generation {
+		req.Header.Set("session-id", "unrelated-client-session")
+	}
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = req
+	c.Params = gin.Params{{Key: "request_id", Value: "task"}}
+	c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{ID: 20, UserID: 10, GroupID: &groupID,
+		Group: &service.Group{ID: groupID, Platform: service.PlatformGrok, AllowImageGeneration: true}, User: &service.User{ID: 10}})
+	c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 10, Concurrency: 5})
+	return c, w
+}
+
+func TestGrokMediaLookupSlotLifecycle(t *testing.T) {
+	for _, scenario := range []string{"normal", "mismatch acquired", "mismatch wait", "full", "queue full", "cancel while waiting", "wait then acquired", "upstream error", "cancel", "panic"} {
+		t.Run(scenario, func(t *testing.T) {
+			h, slots, bindings, upstream := newGrokMediaSlotHandler(t, false, strings.HasPrefix(scenario, "mismatch"))
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			if scenario == "full" || scenario == "mismatch wait" || scenario == "queue full" || scenario == "cancel while waiting" {
+				slots.full = true
+			}
+			if scenario == "queue full" {
+				slots.queueFull = true
+			}
+			if scenario == "cancel while waiting" {
+				slots.onWait = cancel
+			}
+			if scenario == "wait then acquired" {
+				slots.denied = 1
+			}
+			original := upstream.call
+			upstream.call = func(req *http.Request, id int64) (*http.Response, error) {
+				require.Equal(t, int64(1), id)
+				switch scenario {
+				case "upstream error":
+					return nil, errors.New("upstream unavailable")
+				case "cancel":
+					cancel()
+					return nil, context.Canceled
+				case "panic":
+					panic("test upstream panic")
+				}
+				return original(req, id)
+			}
+			for range 20 {
+				c, w := grokMediaSlotContext(ctx, false)
+				h.GrokVideoStatus(c)
+				slots.assertReleased(t)
+				if strings.HasPrefix(scenario, "mismatch") {
+					require.Equal(t, 404, w.Code)
+				}
+				if scenario == "normal" || scenario == "wait then acquired" {
+					require.Equal(t, 200, w.Code)
+				}
+				if scenario == "full" || scenario == "queue full" {
+					require.Equal(t, http.StatusTooManyRequests, w.Code)
+				}
+			}
+			require.Zero(t, bindings.writes, "lookups must preserve owner and TTL")
+			if scenario == "mismatch acquired" {
+				require.Equal(t, 20, slots.released)
+			}
+			if slots.full {
+				require.Zero(t, slots.released)
+				require.Zero(t, upstream.calls)
+			}
+			if scenario == "full" || strings.HasPrefix(scenario, "mismatch") {
+				require.Zero(t, upstream.calls)
+			}
+			switch scenario {
+			case "normal", "wait then acquired", "upstream error", "panic":
+				require.Equal(t, 20, upstream.calls)
+				require.Equal(t, 20, slots.released)
+			case "cancel":
+				require.Equal(t, 1, upstream.calls)
+				require.Equal(t, 1, slots.released)
+			}
+		})
+	}
+}
+
+func TestGrokMediaEligibilityReleasesBeforeSwitch(t *testing.T) {
+	for _, scenario := range []string{"switch", "exhausted", "cancel during probe"} {
+		t.Run(scenario, func(t *testing.T) {
+			h, slots, _, upstream := newGrokMediaSlotHandler(t, true, false)
+			h.maxAccountSwitches = 1
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			probes := 0
+			h.grokMediaEligibilityProber = grokMediaSlotProber(func(context.Context, int64) (bool, string, error) {
+				probes++
+				if scenario == "cancel during probe" {
+					cancel()
+					require.Eventually(t, func() bool { slots.mu.Lock(); defer slots.mu.Unlock(); return len(slots.accounts) == 0 }, time.Second, time.Millisecond)
+				}
+				if scenario == "switch" && probes == 2 {
+					return true, "eligible", nil
+				}
+				return false, "billing_unobserved", errors.New("probe failed")
+			})
+			c, w := grokMediaSlotContext(ctx, true)
+			h.GrokVideoGeneration(c)
+			slots.assertReleased(t)
+			require.Equal(t, 1, slots.peak, "rejected attempts must release before the next selection")
+			if scenario == "switch" {
+				events, _ := c.Get(service.OpsUpstreamErrorsKey)
+				detail, err := json.Marshal(events)
+				require.NoError(t, err)
+				require.Equal(t, 200, w.Code, "body=%s probes=%d upstream=%d events=%s", w.Body.String(), probes, upstream.calls, detail)
+				require.Equal(t, 1, upstream.calls)
+			}
+			if scenario == "exhausted" {
+				require.Equal(t, 503, w.Code)
+				require.Equal(t, 2, slots.released)
+				require.Zero(t, upstream.calls)
+			}
+		})
+	}
+}
+
+func TestGrokMediaVideoLookupOwnerIsolation(t *testing.T) {
+	for _, other := range []string{"user", "api key", "group", "task"} {
+		t.Run(other, func(t *testing.T) {
+			h, slots, bindings, upstream := newGrokMediaSlotHandler(t, false, false)
+			c, w := grokMediaSlotContext(context.Background(), false)
+			key, ok := middleware2.GetAPIKeyFromContext(c)
+			require.True(t, ok)
+			switch other {
+			case "user":
+				c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 11, Concurrency: 5})
+			case "api key":
+				key.ID = 21
+			case "group":
+				groupID := int64(25)
+				key.GroupID = &groupID
+			case "task":
+				c.Params = gin.Params{{Key: "request_id", Value: "other-task"}}
+			}
+			h.GrokVideoStatus(c)
+			require.Equal(t, http.StatusNotFound, w.Code)
+			slots.assertReleased(t)
+			require.Zero(t, slots.acquired)
+			require.Zero(t, upstream.calls)
+			require.Zero(t, bindings.writes)
+		})
+	}
+}
+
+func TestGrokMediaVideoCompletionStillClaimsBillingOnce(t *testing.T) {
+	h, _, bindings, _ := newGrokMediaSlotHandler(t, false, false)
+	c, _ := grokMediaSlotContext(context.Background(), false)
+	key, ok := middleware2.GetAPIKeyFromContext(c)
+	require.True(t, ok)
+	subject := middleware2.AuthSubject{UserID: 10, Concurrency: 5}
+	result := &service.OpenAIForwardResult{ResponseID: "task", Model: "grok-imagine-video",
+		VideoCount: 1, VideoDurationSeconds: 6}
+	for i := range 20 {
+		bill := prepareGrokVideoCompletionBilling(c.Request.Context(), h, zap.NewNop(), key, subject, "task", result)
+		if i == 0 {
+			require.NotNil(t, bill)
+			require.Equal(t, service.StableGrokVideoBillingRequestID("task"), bill.RequestID)
+		} else {
+			require.Nil(t, bill)
+		}
+	}
+	require.Len(t, bindings.billed, 1)
+}