diff --git a/internal/provider/manager.go b/internal/provider/manager.go index 63b890b..d25017d 100644 --- a/internal/provider/manager.go +++ b/internal/provider/manager.go @@ -665,7 +665,10 @@ func (m *Manager) pickNext(ctx context.Context, provider, model string, opts Opt if modelKey != "" && registryRef != nil && !registryRef.ClientSupportsModel(candidate.ID, modelKey) { continue } - candidates = append(candidates, candidate) + // Create a shallow clone for the selector to avoid race conditions. + // The selector needs to read auth.Runtime concurrently with MarkResult updates. + // By cloning here (under RLock), we ensure the selector gets a consistent snapshot. + candidates = append(candidates, candidate.Clone()) } if len(candidates) == 0 { m.mu.RUnlock() diff --git a/internal/provider/quota_group.go b/internal/provider/quota_group.go index 9e141e9..8001abd 100644 --- a/internal/provider/quota_group.go +++ b/internal/provider/quota_group.go @@ -129,8 +129,15 @@ func (rd *AuthRuntimeData) GetProviderData() any { } // getOrCreateQuotaGroupIndex returns the quota group index from auth.Runtime, -// creating it if necessary. This function is thread-safe and preserves -// any existing provider-specific runtime data. +// creating it if necessary. +// +// CRITICAL: This function MUST only be called while holding the Manager's lock (m.mu). +// It is NOT safe to call this concurrently on the same auth object without external +// synchronization. The Manager ensures this by holding m.mu during MarkResult and other +// state-modifying operations. +// +// The returned quotaGroupIndex has its own internal mutex for thread-safe read/write +// operations after initialization. func getOrCreateQuotaGroupIndex(auth *Auth) *quotaGroupIndex { if auth == nil { return nil diff --git a/internal/provider/quota_group_race_test.go b/internal/provider/quota_group_race_test.go new file mode 100644 index 0000000..c3bc970 --- /dev/null +++ b/internal/provider/quota_group_race_test.go @@ -0,0 +1,249 @@ +package provider + +import ( + "context" + "sync" + "testing" + "time" +) + +// TestQuotaGroupIndexManagerProtection tests that quota group index operations +// are properly protected by the Manager's lock, simulating real usage patterns. +func TestQuotaGroupIndexManagerProtection(t *testing.T) { + // Simulate Manager holding lock during state updates + var managerMu sync.RWMutex + + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + Runtime: nil, + } + + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + var wg sync.WaitGroup + iterations := 100 + + // Goroutines simulating MarkResult (write operations with full lock) + for i := 0; i < iterations/2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + + managerMu.Lock() // Manager holds lock during state updates + idx := getOrCreateQuotaGroupIndex(auth) + if idx != nil { + idx.setGroupBlocked("claude", "claude-sonnet", time.Now().Add(time.Minute), time.Now().Add(time.Minute)) + } + managerMu.Unlock() + }() + } + + // Goroutines simulating selector Pick (read operations with read lock) + for i := 0; i < iterations/2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + + time.Sleep(time.Millisecond) // Let some writes happen first + + managerMu.RLock() // Selector gets auth pointer while holding read lock + // In real code, selector receives auth pointer from pickNext + idx := getQuotaGroupIndex(auth) + managerMu.RUnlock() + + // Selector can safely call isGroupBlocked on the index it received + // because quotaGroupIndex has its own internal mutex + if idx != nil { + _, _ = idx.isGroupBlocked("claude", time.Now()) + } + }() + } + + wg.Wait() + + // Verify the index is in a valid state + managerMu.RLock() + idx := getQuotaGroupIndex(auth) + managerMu.RUnlock() + + if idx == nil { + t.Fatal("Index should exist after concurrent operations") + } + + // Should be able to check blocked state without panic + _, _ = idx.isGroupBlocked("claude", time.Now()) +} + +// TestConcurrentQuotaPropagation tests concurrent quota propagation and clearing +// with proper Manager lock protection. +func TestConcurrentQuotaPropagation(t *testing.T) { + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + var managerMu sync.RWMutex + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": {Status: StatusActive}, + "claude-opus-4": {Status: StatusActive}, + "claude-haiku-4": {Status: StatusActive}, + "gemini-2.5-pro": {Status: StatusActive}, + "gemini-2.5-flash": {Status: StatusActive}, + }, + } + + var wg sync.WaitGroup + now := time.Now() + + // Concurrent propagation (with lock, simulating MarkResult) + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + managerMu.Lock() + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: now.Add(time.Minute), + BackoffLevel: 1, + } + _ = propagateQuotaToGroup(auth, "claude-sonnet-4", quota, now.Add(time.Minute), now) + managerMu.Unlock() + }() + } + + // Concurrent clearing (with lock, simulating MarkResult) + for i := 0; i < 50; i++ { + wg.Add(1) + go func() { + defer wg.Done() + time.Sleep(10 * time.Millisecond) + managerMu.Lock() + _ = clearQuotaGroupOnSuccess(auth, "claude-opus-4", now) + managerMu.Unlock() + }() + } + + wg.Wait() + + // Verify auth is in a consistent state + managerMu.RLock() + if auth.ModelStates == nil { + managerMu.RUnlock() + t.Fatal("ModelStates should not be nil") + } + + for model, state := range auth.ModelStates { + if state == nil { + managerMu.RUnlock() + t.Errorf("Model %s has nil state", model) + return + } + } + managerMu.RUnlock() +} + +// TestQuotaGroupIndexConsistency tests that quota index and ModelStates stay in sync. +func TestQuotaGroupIndexConsistency(t *testing.T) { + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": {Status: StatusActive}, + }, + } + + now := time.Now() + retryTime := now.Add(time.Minute) + + // Set quota exceeded via propagation + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: retryTime, + BackoffLevel: 1, + } + propagateQuotaToGroup(auth, "claude-sonnet-4", quota, retryTime, now) + + // Check that isAuthBlockedForModel correctly detects the block for an uninitialized model + // in the same quota group + blocked, _, next := isAuthBlockedForModel(auth, "claude-opus-4", now) + if !blocked { + t.Error("claude-opus-4 should be blocked due to quota group") + } + if next.IsZero() || !next.After(now) { + t.Errorf("Expected future retry time, got %v", next) + } + + // Clear the quota + clearQuotaGroupOnSuccess(auth, "claude-sonnet-4", now) + + // Verify the uninitialized model is no longer blocked + blocked, _, _ = isAuthBlockedForModel(auth, "claude-opus-4", now) + if blocked { + t.Error("claude-opus-4 should not be blocked after quota clear") + } +} + +// TestSelectorWithQuotaGroups tests the full workflow: Manager passes auth clones +// to selector, which safely reads quota state. +func TestSelectorWithQuotaGroups(t *testing.T) { + RegisterQuotaGroupResolver("antigravity", AntigravityQuotaGroupResolver) + + var managerMu sync.RWMutex + auth := &Auth{ + ID: "test-auth", + Provider: "antigravity", + ModelStates: map[string]*ModelState{ + "claude-sonnet-4": {Status: StatusActive}, + }, + } + + selector := &RoundRobinSelector{} + selector.Start() + defer selector.Stop() + + var wg sync.WaitGroup + + // Goroutine 1: Continuously set quota (simulating MarkResult) + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + managerMu.Lock() + quota := QuotaState{ + Exceeded: true, + Reason: "quota", + NextRecoverAt: time.Now().Add(time.Second), + BackoffLevel: 1, + } + propagateQuotaToGroup(auth, "claude-sonnet-4", quota, time.Now().Add(time.Second), time.Now()) + managerMu.Unlock() + time.Sleep(time.Millisecond) + } + }() + + // Goroutine 2: Continuously pick auth (simulating concurrent requests) + // This mimics what pickNext does: clone auth under RLock before passing to selector + wg.Add(1) + go func() { + defer wg.Done() + for i := 0; i < 100; i++ { + managerMu.RLock() + // Manager clones auth before passing to selector (like pickNext does after our fix) + authClone := auth.Clone() + auths := []*Auth{authClone} + managerMu.RUnlock() + + // Selector picks and checks if auth is blocked + _, err := selector.Pick(context.Background(), "antigravity", "claude-opus-4", Options{}, auths) + _ = err // Error is expected when quota is exceeded + time.Sleep(time.Millisecond) + } + }() + + wg.Wait() +}