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.
This commit is contained in:
sususu
2026-08-25 18:55:00 +08:00
parent 998dcfeba2
commit 1f53b2eb03
4 changed files with 286 additions and 58 deletions

View File

@@ -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

View File

@@ -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()

View File

@@ -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
}

View File

@@ -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")
}
}