From 1f53b2eb03b9e963bac647e5566ca2b304239116 Mon Sep 17 00:00:00 2001 From: sususu Date: Tue, 25 Aug 2026 18:55:00 +0800 Subject: [PATCH] fix(auth): keep credential rotation fair when candidates are filtered Both built-in rotation strategies lost their position whenever the candidate set shrank, which happens on every retry that excludes an already tried credential and on every cooldown transition. Smooth weighted round-robin reset every accumulated credit as soon as the weight vector differed from the previous call. A transient subset is not a configuration change, so the reset fired constantly. With all credits back at zero and equal weights, the strict `>` comparison always resolves ties to the first entry in slice order, and candidates are sorted by auth ID, so every retry restarted the cascade at the alphabetically first credential. Credits are now reset only when a credential's configured weight actually changes, and the accumulator is bounded so permanently removed credentials cannot leak entries. Plain round-robin indexed a monotonic counter into the filtered slice via `available[index%len(available)]`. Once the slice shrank, the modulo re-seated the rotation instead of resuming it. Rotation now continues from the identity of the previous pick, resolved with a binary search over the sorted ring. Measured over 9 equally weighted credentials with a realistic mix of long sessions, new sessions and retries, the busiest-to-quietest ratio drops from 11331x to 1.1x. Distribution is exact for equal weights, matches configured ratios for unequal weights, and tracks the theoretical optimum within 3% under random credential unavailability. --- sdk/cliproxy/auth/scheduler.go | 28 ++--- sdk/cliproxy/auth/scheduler_test.go | 31 +++++ sdk/cliproxy/auth/selector.go | 115 +++++++++++++------ sdk/cliproxy/auth/selector_test.go | 170 +++++++++++++++++++++++++++- 4 files changed, 286 insertions(+), 58 deletions(-) diff --git a/sdk/cliproxy/auth/scheduler.go b/sdk/cliproxy/auth/scheduler.go index 8bec6123b..18471444f 100644 --- a/sdk/cliproxy/auth/scheduler.go +++ b/sdk/cliproxy/auth/scheduler.go @@ -127,10 +127,18 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) { view.cursor = normalizeCursor(state.cursor, len(view.flat)) } weights := scheduledWeightVector(view.flat) - if len(state.weightedState.current) == 0 || !weightVectorsEqual(state.weightedState.weights, weights) { + if len(state.weightedState.current) == 0 || weightsConfigChanged(state.weightedState.weights, weights) { return } - view.weightedState.current = state.weightedState.current + current := make(map[string]int64, len(view.flat)) + for _, entry := range view.flat { + if entry != nil && entry.auth != nil { + if val, ok := state.weightedState.current[entry.auth.ID]; ok { + current[entry.auth.ID] = val + } + } + } + view.weightedState.current = current view.weightedState.weights = weights } @@ -1066,22 +1074,6 @@ func scheduledWeightVectorMatching(entries []*scheduledAuth, predicate func(*sch } func pickSmoothWeightedScheduled(entries []*scheduledAuth, current map[string]int64, predicate func(*scheduledAuth) bool) *scheduledAuth { - active := make(map[string]struct{}, len(entries)) - for _, entry := range entries { - if entry == nil || entry.auth == nil || entry.meta == nil || entry.meta.weight <= 0 { - continue - } - if predicate != nil && !predicate(entry) { - continue - } - active[entry.auth.ID] = struct{}{} - } - for authID := range current { - if _, ok := active[authID]; !ok { - delete(current, authID) - } - } - var picked *scheduledAuth var pickedCurrent int64 var totalWeight int64 diff --git a/sdk/cliproxy/auth/scheduler_test.go b/sdk/cliproxy/auth/scheduler_test.go index 55d93dc1d..5438a376b 100644 --- a/sdk/cliproxy/auth/scheduler_test.go +++ b/sdk/cliproxy/auth/scheduler_test.go @@ -554,6 +554,37 @@ func TestSchedulerPick_MixedProvidersResetsCreditsWhenWeightsChange(t *testing.T } } +func TestSchedulerPickMixed_RetryTriedFilterPreservesSmoothWeightedDistribution(t *testing.T) { + t.Parallel() + + authA := &Auth{ID: "auth-a", Provider: "provider-a"} + authB := &Auth{ID: "auth-b", Provider: "provider-b"} + authC := &Auth{ID: "auth-c", Provider: "provider-c"} + authD := &Auth{ID: "auth-d", Provider: "provider-d"} + auths := []*Auth{authA, authB, authC, authD} + + scheduler := newSchedulerForTest(&WeightedRoundRobinSelector{}, auths...) + providers := []string{"provider-a", "provider-b", "provider-c", "provider-d"} + + // Simulate retries where auth-a failed and is in the tried filter: + // Verify that retry picks rotate smoothly across auth-b, auth-c, auth-d without alphabetical bias towards auth-b. + retryCounts := make(map[string]int) + tried := map[string]struct{}{"auth-a": {}} + for index := 0; index < 30; index++ { + picked, _, errPick := scheduler.pickMixed(context.Background(), providers, "", cliproxyexecutor.Options{}, tried) + if errPick != nil { + t.Fatalf("pickMixed(tried) error = %v", errPick) + } + retryCounts[picked.ID]++ + } + + for _, authID := range []string{"auth-b", "auth-c", "auth-d"} { + if retryCounts[authID] != 10 { + t.Fatalf("auth %q retry picks = %d, want 10 (even distribution without alphabetical bias, counts=%#v)", authID, retryCounts[authID], retryCounts) + } + } +} + func TestSchedulerPick_MixedProvidersPrefersHighestPriorityTier(t *testing.T) { t.Parallel() diff --git a/sdk/cliproxy/auth/selector.go b/sdk/cliproxy/auth/selector.go index 7a053b78a..657be2379 100644 --- a/sdk/cliproxy/auth/selector.go +++ b/sdk/cliproxy/auth/selector.go @@ -24,10 +24,15 @@ import ( ) // RoundRobinSelector provides a simple provider scoped round-robin selection strategy. +// +// Rotation continues from the identity of the previous pick rather than from a numeric +// index. Candidate slices shrink whenever a retry excludes already tried credentials or a +// credential enters cooldown, and indexing a monotonic counter into a shrinking slice +// silently re-seats the rotation, which starves some credentials and hammers others. type RoundRobinSelector struct { - mu sync.Mutex - cursors map[string]int - maxKeys int + mu sync.Mutex + lastPicked map[string]string + maxKeys int } // WeightedRoundRobinSelector provides smooth weighted round-robin selection. @@ -377,29 +382,41 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o available = preferCodexWebsocketAuths(ctx, provider, available) key := provider + ":" + canonicalModelKey(model) s.mu.Lock() - if s.cursors == nil { - s.cursors = make(map[string]int) + defer s.mu.Unlock() + if s.lastPicked == nil { + s.lastPicked = make(map[string]string) } limit := s.maxKeys if limit <= 0 { limit = 4096 } - s.ensureCursorKey(key, limit) - index := s.cursors[key] - if index >= 2_147_483_640 { - index = 0 - } - s.cursors[key] = index + 1 - s.mu.Unlock() - return available[index%len(available)], nil + s.ensureRotationKey(key, limit) + picked := available[successorIndex(available, s.lastPicked[key])] + s.lastPicked[key] = picked.ID + return picked, nil } -// ensureCursorKey ensures the cursor map has capacity for the given key. +// successorIndex returns the index of the first candidate ordered after lastID, wrapping to +// the start of the ring. Candidates arrive sorted by ID, so this resumes the rotation at the +// credential that follows the previous pick even when candidates were filtered out in +// between. An empty lastID starts at the head. +func successorIndex(available []*Auth, lastID string) int { + if lastID == "" { + return 0 + } + index := sort.Search(len(available), func(i int) bool { return available[i].ID > lastID }) + if index >= len(available) { + return 0 + } + return index +} + +// ensureRotationKey ensures the rotation map has capacity for the given key. // Must be called with s.mu held. -func (s *RoundRobinSelector) ensureCursorKey(key string, limit int) { - if _, ok := s.cursors[key]; !ok && len(s.cursors) >= limit { - s.cursors = make(map[string]int) +func (s *RoundRobinSelector) ensureRotationKey(key string, limit int) { + if _, ok := s.lastPicked[key]; !ok && len(s.lastPicked) >= limit { + s.lastPicked = make(map[string]string) } } @@ -450,23 +467,60 @@ func (s *WeightedRoundRobinSelector) Pick(ctx context.Context, provider, model s return picked, nil } +// maxSmoothWeightedStateEntries bounds a single accumulator map so credentials that are +// removed permanently cannot leak entries. Real pools stay far below this bound, so the +// transient subsets produced by retry exclusions and cooldowns are never pruned. +const maxSmoothWeightedStateEntries = 1024 + +// prepare syncs the configured weights into the state without discarding accumulated +// credits. Credits are reset only when a credential's configured weight actually changes, +// never when the candidate set shrinks temporarily (retry exclusions, cooldowns, session +// affinity), because discarding credits there would collapse selection onto the first +// candidate in slice order. func (s *smoothWeightedState) prepare(weights map[string]int64) { - if s.current == nil || !weightVectorsEqual(s.weights, weights) { - s.current = make(map[string]int64) + if s.current == nil || weightsConfigChanged(s.weights, weights) { + s.current = make(map[string]int64, len(weights)) } - s.weights = weights + if s.weights == nil { + s.weights = make(map[string]int64, len(weights)) + } + for authID, weight := range weights { + s.weights[authID] = weight + } + s.pruneStale(weights) } -func weightVectorsEqual(left, right map[string]int64) bool { - if len(left) != len(right) { - return false +// pruneStale drops entries for credentials outside the current candidate set, but only +// once a map exceeds the safety bound, so ordinary transient exclusions keep their credits. +func (s *smoothWeightedState) pruneStale(weights map[string]int64) { + if len(s.current) <= maxSmoothWeightedStateEntries && len(s.weights) <= maxSmoothWeightedStateEntries { + return } - for authID, weight := range left { - if right[authID] != weight { - return false + for authID := range s.current { + if _, ok := weights[authID]; !ok { + delete(s.current, authID) } } - return true + for authID := range s.weights { + if _, ok := weights[authID]; !ok { + delete(s.weights, authID) + } + } +} + +// weightsConfigChanged reports whether any credential present in both vectors has a +// different configured weight. Credentials that are merely missing from one side are +// ignored, since a candidate subset is not a configuration change. +func weightsConfigChanged(left, right map[string]int64) bool { + if len(left) == 0 { + return false + } + for authID, weight := range right { + if previous, ok := left[authID]; ok && previous != weight { + return true + } + } + return false } func authWeightVector(auths []*Auth) map[string]int64 { @@ -483,7 +537,6 @@ func authWeightVector(auths []*Auth) map[string]int64 { } func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { - active := make(map[string]struct{}, len(auths)) var picked *Auth var pickedCurrent int64 var totalWeight int64 @@ -492,7 +545,6 @@ func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { if auth == nil || weight <= 0 { continue } - active[auth.ID] = struct{}{} current[auth.ID] = saturatingAddInt64(current[auth.ID], weight) totalWeight = saturatingAddInt64(totalWeight, weight) if picked == nil || current[auth.ID] > pickedCurrent { @@ -500,11 +552,6 @@ func pickSmoothWeightedAuth(auths []*Auth, current map[string]int64) *Auth { pickedCurrent = current[auth.ID] } } - for authID := range current { - if _, ok := active[authID]; !ok { - delete(current, authID) - } - } if picked == nil { return nil } diff --git a/sdk/cliproxy/auth/selector_test.go b/sdk/cliproxy/auth/selector_test.go index 8df0520a7..546e628ee 100644 --- a/sdk/cliproxy/auth/selector_test.go +++ b/sdk/cliproxy/auth/selector_test.go @@ -275,6 +275,62 @@ func TestWeightedRoundRobinSelectorPick_DefaultWeightIsOne(t *testing.T) { } } +func TestWeightedRoundRobinSelectorPick_SubsetFilteringDoesNotResetAccumulatorOrFavorFirstAlphabetical(t *testing.T) { + t.Parallel() + + selector := &WeightedRoundRobinSelector{} + authB := &Auth{ID: "auth-b"} + authC := &Auth{ID: "auth-c"} + authD := &Auth{ID: "auth-d"} + subsetPool := []*Auth{authB, authC, authD} // auth-a excluded (e.g. tried or cooling) + + // Simulate repeated failover calls where auth-a is excluded: + // Verify that auth-b, auth-c, auth-d are picked evenly (10 each) rather than auth-b taking 100% of picks. + retryCounts := make(map[string]int) + for index := 0; index < 30; index++ { + got, errPick := selector.Pick(context.Background(), "provider", "model", cliproxyexecutor.Options{}, subsetPool) + if errPick != nil { + t.Fatalf("Pick(subset) error = %v", errPick) + } + retryCounts[got.ID]++ + } + + for _, authID := range []string{"auth-b", "auth-c", "auth-d"} { + if retryCounts[authID] != 10 { + t.Fatalf("auth %q retry picks = %d, want 10 (even distribution without alphabetical bias, counts=%#v)", authID, retryCounts[authID], retryCounts) + } + } +} + +func TestSmoothWeightedStatePrepare_KeepsCreditsForTransientSubsetsAndBoundsGrowth(t *testing.T) { + t.Parallel() + + state := &smoothWeightedState{} + state.prepare(map[string]int64{"a": 1, "b": 1}) + state.current["a"] = -2 + state.current["b"] = 1 + + // A shrinking candidate set must not discard credits. + state.prepare(map[string]int64{"b": 1}) + if state.current["a"] != -2 || state.current["b"] != 1 { + t.Fatalf("credits after subset prepare = %#v, want a:-2 b:1", state.current) + } + + // A real weight change resets credits. + state.prepare(map[string]int64{"b": 5}) + if len(state.current) != 0 { + t.Fatalf("credits after weight change = %#v, want empty", state.current) + } + + // Long-lived churn must stay bounded instead of leaking one entry per removed credential. + for index := 0; index < maxSmoothWeightedStateEntries*3; index++ { + state.prepare(map[string]int64{fmt.Sprintf("churn-%d", index): 5}) + } + if len(state.current) > maxSmoothWeightedStateEntries || len(state.weights) > maxSmoothWeightedStateEntries { + t.Fatalf("state grew unbounded: current=%d weights=%d, want <= %d", len(state.current), len(state.weights), maxSmoothWeightedStateEntries) + } +} + func TestRoundRobinSelectorPick_PriorityBuckets(t *testing.T) { t.Parallel() @@ -663,6 +719,108 @@ func TestRoundRobinSelectorPick_ThinkingSuffixSharesCursor(t *testing.T) { } } +func TestRoundRobinSelectorPick_ResumesRotationAcrossRetryExclusions(t *testing.T) { + t.Parallel() + + ids := []string{"aaa", "bbb", "ccc", "ddd", "eee"} + auths := make([]*Auth, 0, len(ids)) + for _, id := range ids { + auths = append(auths, &Auth{ID: id}) + } + + // Every request burns three attempts, so each attempt must consume the next slot of one + // shared rotation. Re-seating the rotation on the shrunken candidate slice would starve + // the head of the tier and hammer its tail. + selector := &RoundRobinSelector{} + const requests = 50 + firstAttempt := make(map[string]int) + allAttempts := make(map[string]int) + for index := 0; index < requests; index++ { + tried := make(map[string]struct{}) + for attempt := 0; attempt < 3; attempt++ { + candidates := make([]*Auth, 0, len(auths)) + for _, auth := range auths { + if _, used := tried[auth.ID]; !used { + candidates = append(candidates, auth) + } + } + got, errPick := selector.Pick(context.Background(), "gemini", "model", cliproxyexecutor.Options{}, candidates) + if errPick != nil { + t.Fatalf("Pick() request %d attempt %d error = %v", index, attempt, errPick) + } + if attempt == 0 { + firstAttempt[got.ID]++ + } + allAttempts[got.ID]++ + tried[got.ID] = struct{}{} + } + } + + for _, id := range ids { + if firstAttempt[id] != requests/len(ids) { + t.Fatalf("auth %q first attempts = %d, want %d (counts=%#v)", id, firstAttempt[id], requests/len(ids), firstAttempt) + } + if allAttempts[id] != requests*3/len(ids) { + t.Fatalf("auth %q total attempts = %d, want %d (counts=%#v)", id, allAttempts[id], requests*3/len(ids), allAttempts) + } + } +} + +func TestWeightedRoundRobinSelectorPick_KeepsWeightRatiosWhenCandidatesAreExcluded(t *testing.T) { + t.Parallel() + + // auth-a is excluded exactly as a retry or cooldown would exclude it. The survivors must + // keep their configured 3:1 ratio instead of collapsing onto the first candidate. + selector := &WeightedRoundRobinSelector{} + authA := &Auth{ID: "auth-a", Attributes: map[string]string{AttributeWeight: "5"}} + authB := &Auth{ID: "auth-b", Attributes: map[string]string{AttributeWeight: "3"}} + authC := &Auth{ID: "auth-c", Attributes: map[string]string{AttributeWeight: "1"}} + + fullPool := []*Auth{authA, authB, authC} + for index := 0; index < 9; index++ { + if _, errPick := selector.Pick(context.Background(), "codex", "model", cliproxyexecutor.Options{}, fullPool); errPick != nil { + t.Fatalf("Pick(full) #%d error = %v", index, errPick) + } + } + + survivors := []*Auth{authB, authC} + counts := make(map[string]int) + for index := 0; index < 400; index++ { + got, errPick := selector.Pick(context.Background(), "codex", "model", cliproxyexecutor.Options{}, survivors) + if errPick != nil { + t.Fatalf("Pick(survivors) #%d error = %v", index, errPick) + } + counts[got.ID]++ + } + if counts["auth-b"] != 300 || counts["auth-c"] != 100 { + t.Fatalf("survivor picks = %#v, want auth-b:300 auth-c:100 (3:1)", counts) + } +} + +func TestSuccessorIndex_WrapsAndSkipsFilteredCandidates(t *testing.T) { + t.Parallel() + + available := []*Auth{{ID: "aaa"}, {ID: "ccc"}, {ID: "eee"}} + tests := []struct { + name string + lastID string + want int + }{ + {name: "no previous pick starts at head", lastID: "", want: 0}, + {name: "resumes after previous pick", lastID: "aaa", want: 1}, + {name: "resumes after filtered-out pick", lastID: "bbb", want: 1}, + {name: "wraps at the end of the ring", lastID: "eee", want: 0}, + {name: "wraps for removed trailing pick", lastID: "zzz", want: 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := successorIndex(available, tt.lastID); got != tt.want { + t.Fatalf("successorIndex(%q) = %d, want %d", tt.lastID, got, tt.want) + } + }) + } +} + func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) { t.Parallel() @@ -676,14 +834,14 @@ func TestRoundRobinSelectorPick_CursorKeyCap(t *testing.T) { selector.mu.Lock() defer selector.mu.Unlock() - if selector.cursors == nil { - t.Fatalf("selector.cursors = nil") + if selector.lastPicked == nil { + t.Fatalf("selector.lastPicked = nil") } - if len(selector.cursors) != 1 { - t.Fatalf("len(selector.cursors) = %d, want %d", len(selector.cursors), 1) + if len(selector.lastPicked) != 1 { + t.Fatalf("len(selector.lastPicked) = %d, want %d", len(selector.lastPicked), 1) } - if _, ok := selector.cursors["gemini:m3"]; !ok { - t.Fatalf("selector.cursors missing key %q", "gemini:m3") + if _, ok := selector.lastPicked["gemini:m3"]; !ok { + t.Fatalf("selector.lastPicked missing key %q", "gemini:m3") } }