mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-29 16:20:27 +00:00
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).
This commit is contained in:
@@ -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)
|
||||
}
|
||||
@@ -332,6 +332,36 @@ func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount(
|
||||
return s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey)
|
||||
}
|
||||
|
||||
// SelectGrokMediaVideoRequestAccount only admits the already authenticated
|
||||
// task owner. Generic sticky fallback can query another account and overwrite
|
||||
// the ownership key; video lookups must neither escape nor refresh that key.
|
||||
func (s *OpenAIGatewayService) SelectGrokMediaVideoRequestAccount(
|
||||
ctx context.Context, groupID *int64, sessionHash string, accountID int64, requestedModel string,
|
||||
) (*AccountSelectionResult, OpenAIAccountScheduleDecision, error) {
|
||||
decision := OpenAIAccountScheduleDecision{Layer: openAIAccountScheduleLayerSessionSticky}
|
||||
if accountID <= 0 || strings.TrimSpace(sessionHash) == "" {
|
||||
return nil, decision, ErrNoAvailableAccounts
|
||||
}
|
||||
ctx = s.withOpenAIGroupPrivacyRequirement(WithOpenAIProfitControlSuppressed(ctx), groupID)
|
||||
scheduler := &defaultOpenAIAccountScheduler{service: s}
|
||||
selection, _, err := scheduler.selectBySessionHash(ctx, OpenAIAccountScheduleRequest{
|
||||
GroupID: groupID, Platform: PlatformGrok, SessionHash: sessionHash,
|
||||
StickyAccountID: accountID, PreserveStickyBinding: true, DisableStickyEscape: true,
|
||||
RequestedModel: requestedModel, RequiredTransport: OpenAIUpstreamTransportHTTPSSE,
|
||||
RequirePrivacySet: s.openAIGroupRequiresPrivacySet(ctx, groupID),
|
||||
})
|
||||
if err != nil {
|
||||
return nil, decision, err
|
||||
}
|
||||
if selection == nil || selection.Account == nil {
|
||||
return nil, decision, ErrNoAvailableAccounts
|
||||
}
|
||||
decision.StickySessionHit = true
|
||||
decision.SelectedAccountID = selection.Account.ID
|
||||
decision.SelectedAccountType = selection.Account.Type
|
||||
return selection, decision, nil
|
||||
}
|
||||
|
||||
// GrokVideoPendingBilling is the create-time snapshot used when status polling
|
||||
// first observes a completed video URL. Status may omit model/duration; we fall
|
||||
// back to this snapshot, then defaults.
|
||||
|
||||
@@ -0,0 +1,110 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestSelectGrokMediaVideoRequestAccountPreservesOwner(t *testing.T) {
|
||||
for _, state := range []string{"available", "full", "unavailable", "wrong group", "missing", "invalid id"} {
|
||||
t.Run(state, func(t *testing.T) {
|
||||
groupID := int64(24)
|
||||
ownerID := int64(1)
|
||||
owner := Account{ID: ownerID, Platform: PlatformGrok, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 50, GroupIDs: []int64{groupID}}
|
||||
other := owner
|
||||
other.ID = 2
|
||||
if state == "unavailable" {
|
||||
until := time.Now().Add(time.Minute)
|
||||
owner.TempUnschedulableUntil = &until
|
||||
}
|
||||
if state == "wrong group" {
|
||||
owner.GroupIDs = []int64{25}
|
||||
}
|
||||
accounts := []Account{owner, other}
|
||||
if state == "missing" {
|
||||
accounts = []Account{other}
|
||||
}
|
||||
if state == "invalid id" {
|
||||
ownerID = 0
|
||||
}
|
||||
var acquired, released []int64
|
||||
cache := &schedulerTestGatewayCache{}
|
||||
cfg := &config.Config{}
|
||||
cfg.Gateway.OpenAIScheduler.StickyEscapeEnabled = true
|
||||
cfg.Gateway.Scheduling.StickySessionWaitTimeout = time.Second
|
||||
cfg.Gateway.Scheduling.StickySessionMaxWaiting = 3
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: schedulerTestOpenAIAccountRepo{accounts: accounts}, cache: cache, cfg: cfg,
|
||||
concurrencyService: NewConcurrencyService(schedulerTestConcurrencyCache{
|
||||
acquireResults: map[int64]bool{1: state != "full", 2: true},
|
||||
acquiredIDs: &acquired, releasedIDs: &released,
|
||||
}),
|
||||
}
|
||||
ctx := context.Background()
|
||||
require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "task", 10, 20, 1))
|
||||
sessionHash := GrokMediaVideoRequestSessionHash("task", 10, 20)
|
||||
for range 20 {
|
||||
selection, decision, err := svc.SelectGrokMediaVideoRequestAccount(ctx, &groupID, sessionHash, ownerID, "")
|
||||
switch state {
|
||||
case "available":
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), selection.Account.ID)
|
||||
require.True(t, selection.Acquired)
|
||||
require.True(t, decision.StickySessionHit)
|
||||
selection.ReleaseFunc()
|
||||
case "full":
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), selection.Account.ID)
|
||||
require.False(t, selection.Acquired)
|
||||
require.Nil(t, selection.ReleaseFunc)
|
||||
require.Equal(t, int64(1), selection.WaitPlan.AccountID)
|
||||
default:
|
||||
require.ErrorIs(t, err, ErrNoAvailableAccounts)
|
||||
require.Nil(t, selection)
|
||||
}
|
||||
bound, err := svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "task", 10, 20)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(1), bound)
|
||||
}
|
||||
require.NotContains(t, acquired, int64(2))
|
||||
require.Empty(t, cache.deletedSessions)
|
||||
if state == "available" {
|
||||
require.Len(t, released, 20)
|
||||
} else {
|
||||
require.Empty(t, released)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokVideoStickySelectionIgnoresHealthEscape(t *testing.T) {
|
||||
groupID := int64(24)
|
||||
account := Account{ID: 1, Platform: PlatformGrok, Type: AccountTypeAPIKey,
|
||||
Status: StatusActive, Schedulable: true, Concurrency: 50, GroupIDs: []int64{groupID}}
|
||||
svc := &OpenAIGatewayService{
|
||||
accountRepo: schedulerTestOpenAIAccountRepo{accounts: []Account{account}},
|
||||
cache: &schedulerTestGatewayCache{},
|
||||
}
|
||||
stats := newOpenAIAccountRuntimeStats()
|
||||
for range 20 {
|
||||
stats.report(1, false, nil)
|
||||
}
|
||||
scheduler := &defaultOpenAIAccountScheduler{service: svc, stats: stats}
|
||||
req := OpenAIAccountScheduleRequest{GroupID: &groupID, Platform: PlatformGrok,
|
||||
SessionHash: "task", StickyAccountID: 1, PreserveStickyBinding: true}
|
||||
selection, escaped, err := scheduler.selectBySessionHash(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.Nil(t, selection)
|
||||
require.True(t, escaped)
|
||||
req.DisableStickyEscape = true
|
||||
selection, escaped, err = scheduler.selectBySessionHash(context.Background(), req)
|
||||
require.NoError(t, err)
|
||||
require.False(t, escaped)
|
||||
require.True(t, selection.Acquired)
|
||||
selection.ReleaseFunc()
|
||||
}
|
||||
@@ -77,6 +77,9 @@ type OpenAIAccountScheduleRequest struct {
|
||||
StickyWeighted bool
|
||||
SubscriptionPriority bool
|
||||
PreserveStickyBinding bool
|
||||
// DisableStickyEscape keeps task-owner lookups on their account even when
|
||||
// generic sticky health or concurrency heuristics would prefer another.
|
||||
DisableStickyEscape bool
|
||||
RequirePrivacySet bool
|
||||
PreviousResponseID string
|
||||
PreviousResponseCanMove bool
|
||||
@@ -555,7 +558,7 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
return nil, false, nil
|
||||
}
|
||||
escapeCfg := s.service.openAIStickyEscapeConfig()
|
||||
if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape {
|
||||
if reason, errorRate, ttft, shouldEscape := s.shouldEscapeStickyAccount(accountID, escapeCfg); shouldEscape && !req.DisableStickyEscape {
|
||||
slog.Info("sticky_escape_triggered",
|
||||
"account_id", accountID,
|
||||
"reason", reason,
|
||||
@@ -565,6 +568,9 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
return nil, true, nil
|
||||
}
|
||||
result, acquireErr := s.service.tryAcquireAccountSlot(ctx, accountID, account.Concurrency)
|
||||
if acquireErr != nil && req.DisableStickyEscape {
|
||||
return nil, false, acquireErr
|
||||
}
|
||||
if acquireErr == nil && result != nil && result.Acquired {
|
||||
if !req.PreserveStickyBinding {
|
||||
_ = s.service.refreshStickySessionTTL(ctx, req.GroupID, sessionHash, s.service.openAIWSSessionStickyTTL())
|
||||
@@ -579,7 +585,7 @@ func (s *defaultOpenAIAccountScheduler) selectBySessionHash(
|
||||
cfg := s.service.schedulingConfig()
|
||||
// WaitPlan.MaxConcurrency 使用 Concurrency(非 EffectiveLoadFactor),因为 WaitPlan 控制的是 Redis 实际并发槽位等待。
|
||||
if s.service.concurrencyService != nil {
|
||||
if escapeCfg.enabled && acquireErr == nil && result != nil && !result.Acquired {
|
||||
if escapeCfg.enabled && !req.DisableStickyEscape && acquireErr == nil && result != nil && !result.Acquired {
|
||||
errorRate, ttft, _ := s.stats.snapshot(accountID)
|
||||
slog.Info("sticky_escape_triggered",
|
||||
"account_id", accountID,
|
||||
|
||||
Reference in New Issue
Block a user