mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-03 06:35:00 +08:00
fix(auth): preserve round-robin successor across ready view rebuilds
Keep round-robin rotation stable in the scheduler fast path when credentials enter cooldown or are removed. Replace the numeric readyView cursor modulo normalization with ID-based successor binary search, aligning readyView with RoundRobinSelector.
This commit is contained in:
sususu
@@ -85,7 +85,7 @@ type readyBucket struct {
|
||||
// readyView holds the selection order for flat round-robin traversal.
|
||||
type readyView struct {
|
||||
flat []*scheduledAuth
|
||||
cursor int
|
||||
lastPicked string
|
||||
weightedState smoothWeightedState
|
||||
}
|
||||
|
||||
@@ -93,7 +93,7 @@ type readyView struct {
|
||||
type cooldownQueue []*scheduledAuth
|
||||
|
||||
type readyViewCursorState struct {
|
||||
cursor int
|
||||
lastPicked string
|
||||
weightedState smoothWeightedState
|
||||
}
|
||||
|
||||
@@ -103,7 +103,7 @@ type readyBucketCursorState struct {
|
||||
}
|
||||
|
||||
func snapshotReadyViewCursors(view readyView) readyViewCursorState {
|
||||
state := readyViewCursorState{cursor: view.cursor}
|
||||
state := readyViewCursorState{lastPicked: view.lastPicked}
|
||||
if len(view.weightedState.current) > 0 {
|
||||
state.weightedState.current = make(map[string]int64, len(view.weightedState.current))
|
||||
for authID, current := range view.weightedState.current {
|
||||
@@ -123,9 +123,7 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) {
|
||||
if view == nil {
|
||||
return
|
||||
}
|
||||
if len(view.flat) > 0 {
|
||||
view.cursor = normalizeCursor(state.cursor, len(view.flat))
|
||||
}
|
||||
view.lastPicked = state.lastPicked
|
||||
weights := scheduledWeightVector(view.flat)
|
||||
if len(state.weightedState.current) == 0 || weightsConfigChanged(state.weightedState.weights, weights) {
|
||||
return
|
||||
@@ -142,17 +140,6 @@ func restoreReadyViewCursors(view *readyView, state readyViewCursorState) {
|
||||
view.weightedState.weights = weights
|
||||
}
|
||||
|
||||
func normalizeCursor(cursor, size int) int {
|
||||
if size <= 0 || cursor <= 0 {
|
||||
return 0
|
||||
}
|
||||
cursor = cursor % size
|
||||
if cursor < 0 {
|
||||
cursor += size
|
||||
}
|
||||
return cursor
|
||||
}
|
||||
|
||||
// newAuthScheduler constructs an empty scheduler configured for the supplied selector strategy.
|
||||
func newAuthScheduler(selector Selector) *authScheduler {
|
||||
return &authScheduler{
|
||||
@@ -1030,22 +1017,37 @@ func (v *readyView) pickRoundRobin(predicate func(*scheduledAuth) bool) *schedul
|
||||
if len(v.flat) == 0 {
|
||||
return nil
|
||||
}
|
||||
start := 0
|
||||
if len(v.flat) > 0 {
|
||||
start = v.cursor % len(v.flat)
|
||||
}
|
||||
start := scheduledSuccessorIndex(v.flat, v.lastPicked)
|
||||
for offset := 0; offset < len(v.flat); offset++ {
|
||||
index := (start + offset) % len(v.flat)
|
||||
entry := v.flat[index]
|
||||
if entry == nil || entry.auth == nil {
|
||||
continue
|
||||
}
|
||||
if predicate != nil && !predicate(entry) {
|
||||
continue
|
||||
}
|
||||
v.cursor = index + 1
|
||||
v.lastPicked = entry.auth.ID
|
||||
return entry
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// scheduledSuccessorIndex returns the index of the first scheduled candidate ordered after
|
||||
// lastID, wrapping to the start of the ring. Candidates in readyView arrive sorted by auth ID.
|
||||
func scheduledSuccessorIndex(entries []*scheduledAuth, lastID string) int {
|
||||
if lastID == "" {
|
||||
return 0
|
||||
}
|
||||
index := sort.Search(len(entries), func(i int) bool {
|
||||
return entries[i].auth.ID > lastID
|
||||
})
|
||||
if index >= len(entries) {
|
||||
return 0
|
||||
}
|
||||
return index
|
||||
}
|
||||
|
||||
// pickWeighted returns the next ready entry using smooth weighted round-robin.
|
||||
func (v *readyView) pickWeighted(predicate func(*scheduledAuth) bool) *scheduledAuth {
|
||||
if v == nil || len(v.flat) == 0 {
|
||||
|
||||
@@ -585,6 +585,254 @@ func TestSchedulerPickMixed_RetryTriedFilterPreservesSmoothWeightedDistribution(
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadyViewRoundRobinPreservesSuccessorAcrossRebuild(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
entry := func(id string) *scheduledAuth {
|
||||
return &scheduledAuth{auth: &Auth{ID: id}}
|
||||
}
|
||||
|
||||
t.Run("a cooling resumes at b", func(t *testing.T) {
|
||||
original := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B"), entry("C")},
|
||||
}
|
||||
if got := original.pickRoundRobin(nil); got == nil || got.auth.ID != "A" {
|
||||
t.Fatalf("first pick = %v, want A", got)
|
||||
}
|
||||
|
||||
state := snapshotReadyViewCursors(original)
|
||||
rebuilt := readyView{
|
||||
flat: []*scheduledAuth{entry("B"), entry("C")},
|
||||
}
|
||||
restoreReadyViewCursors(&rebuilt, state)
|
||||
|
||||
got := rebuilt.pickRoundRobin(nil)
|
||||
if got == nil || got.auth.ID != "B" {
|
||||
t.Fatalf("pick after A cooldown = %v, want B", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("b cooling resumes at c", func(t *testing.T) {
|
||||
original := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B"), entry("C")},
|
||||
}
|
||||
// Pick A, then B
|
||||
if got := original.pickRoundRobin(nil); got == nil || got.auth.ID != "A" {
|
||||
t.Fatalf("first pick = %v, want A", got)
|
||||
}
|
||||
if got := original.pickRoundRobin(nil); got == nil || got.auth.ID != "B" {
|
||||
t.Fatalf("second pick = %v, want B", got)
|
||||
}
|
||||
|
||||
state := snapshotReadyViewCursors(original)
|
||||
rebuilt := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("C")},
|
||||
}
|
||||
restoreReadyViewCursors(&rebuilt, state)
|
||||
|
||||
got := rebuilt.pickRoundRobin(nil)
|
||||
if got == nil || got.auth.ID != "C" {
|
||||
t.Fatalf("pick after B cooldown = %v, want C", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("c cooling wraps to a", func(t *testing.T) {
|
||||
original := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B"), entry("C")},
|
||||
}
|
||||
// Pick A, B, C
|
||||
for _, want := range []string{"A", "B", "C"} {
|
||||
if got := original.pickRoundRobin(nil); got == nil || got.auth.ID != want {
|
||||
t.Fatalf("pick = %v, want %s", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
state := snapshotReadyViewCursors(original)
|
||||
rebuilt := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B")},
|
||||
}
|
||||
restoreReadyViewCursors(&rebuilt, state)
|
||||
|
||||
got := rebuilt.pickRoundRobin(nil)
|
||||
if got == nil || got.auth.ID != "A" {
|
||||
t.Fatalf("pick after C cooldown = %v, want A", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("recovery preserves successor", func(t *testing.T) {
|
||||
original := readyView{
|
||||
flat: []*scheduledAuth{entry("B"), entry("C")},
|
||||
}
|
||||
if got := original.pickRoundRobin(nil); got == nil || got.auth.ID != "B" {
|
||||
t.Fatalf("first pick = %v, want B", got)
|
||||
}
|
||||
|
||||
state := snapshotReadyViewCursors(original)
|
||||
// A recovered and is prepended back
|
||||
rebuilt := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B"), entry("C")},
|
||||
}
|
||||
restoreReadyViewCursors(&rebuilt, state)
|
||||
|
||||
got := rebuilt.pickRoundRobin(nil)
|
||||
if got == nil || got.auth.ID != "C" {
|
||||
t.Fatalf("pick after A recovery = %v, want C", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("retry exclusion resumes without rebuild", func(t *testing.T) {
|
||||
view := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B"), entry("C")},
|
||||
}
|
||||
if got := view.pickRoundRobin(nil); got == nil || got.auth.ID != "A" {
|
||||
t.Fatalf("first pick = %v, want A", got)
|
||||
}
|
||||
got := view.pickRoundRobin(func(candidate *scheduledAuth) bool {
|
||||
return candidate.auth.ID != "B"
|
||||
})
|
||||
if got == nil || got.auth.ID != "C" {
|
||||
t.Fatalf("pick after excluding B = %v, want C", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("multiple cooldown skips to first surviving successor", func(t *testing.T) {
|
||||
original := readyView{
|
||||
flat: []*scheduledAuth{entry("A"), entry("B"), entry("C"), entry("D")},
|
||||
}
|
||||
if got := original.pickRoundRobin(nil); got == nil || got.auth.ID != "A" {
|
||||
t.Fatalf("first pick = %v, want A", got)
|
||||
}
|
||||
|
||||
state := snapshotReadyViewCursors(original)
|
||||
rebuilt := readyView{
|
||||
flat: []*scheduledAuth{entry("C"), entry("D")},
|
||||
}
|
||||
restoreReadyViewCursors(&rebuilt, state)
|
||||
|
||||
got := rebuilt.pickRoundRobin(nil)
|
||||
if got == nil || got.auth.ID != "C" {
|
||||
t.Fatalf("pick after A and B cooldown = %v, want C", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestScheduledSuccessorIndex_WrapsAndSkipsFilteredCandidates(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
entries := []*scheduledAuth{
|
||||
{auth: &Auth{ID: "aaa"}},
|
||||
{auth: &Auth{ID: "ccc"}},
|
||||
{auth: &Auth{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 := scheduledSuccessorIndex(entries, tt.lastID); got != tt.want {
|
||||
t.Fatalf("scheduledSuccessorIndex(%q) = %d, want %d", tt.lastID, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
if got := scheduledSuccessorIndex(nil, "aaa"); got != 0 {
|
||||
t.Fatalf("scheduledSuccessorIndex(nil, aaa) = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerRoundRobinPreservesSuccessorAcrossCooldown(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
model := "test-successor-model"
|
||||
authIDs := []string{"successor-auth-a", "successor-auth-b", "successor-auth-c"}
|
||||
registerSchedulerModels(t, "gemini", model, authIDs...)
|
||||
|
||||
for _, id := range authIDs {
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: id, Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(%s) error = %v", id, errRegister)
|
||||
}
|
||||
}
|
||||
|
||||
got, errPick := manager.scheduler.pickSingle(context.Background(), "gemini", model, cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickSingle #1 error = %v", errPick)
|
||||
}
|
||||
if got == nil || got.ID != "successor-auth-a" {
|
||||
t.Fatalf("pickSingle #1 = %v, want successor-auth-a", got)
|
||||
}
|
||||
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: "successor-auth-a",
|
||||
Provider: "gemini",
|
||||
Model: model,
|
||||
Success: false,
|
||||
Error: &Error{HTTPStatus: 429, Message: "rate limit"},
|
||||
})
|
||||
|
||||
got, errPick = manager.scheduler.pickSingle(context.Background(), "gemini", model, cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickSingle #2 after successor-auth-a cooldown error = %v", errPick)
|
||||
}
|
||||
if got == nil || got.ID != "successor-auth-b" {
|
||||
t.Fatalf("pickSingle #2 after successor-auth-a cooldown = %v, want successor-auth-b", got)
|
||||
}
|
||||
|
||||
manager.MarkResult(context.Background(), Result{
|
||||
AuthID: "successor-auth-b",
|
||||
Provider: "gemini",
|
||||
Model: model,
|
||||
Success: false,
|
||||
Error: &Error{HTTPStatus: 429, Message: "rate limit"},
|
||||
})
|
||||
|
||||
got, errPick = manager.scheduler.pickSingle(context.Background(), "gemini", model, cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickSingle #3 after successor-auth-b cooldown error = %v", errPick)
|
||||
}
|
||||
if got == nil || got.ID != "successor-auth-c" {
|
||||
t.Fatalf("pickSingle #3 after successor-auth-b cooldown = %v, want successor-auth-c", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerPick_RoundRobinPreservesWebsocketSuccessorAcrossCooldown(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
wsA := &Auth{ID: "codex-ws-a", Provider: "codex", Attributes: map[string]string{"websockets": "true"}}
|
||||
wsB := &Auth{ID: "codex-ws-b", Provider: "codex", Attributes: map[string]string{"websockets": "true"}}
|
||||
wsC := &Auth{ID: "codex-ws-c", Provider: "codex", Attributes: map[string]string{"websockets": "true"}}
|
||||
httpOnly := &Auth{ID: "codex-http", Provider: "codex"}
|
||||
scheduler := newSchedulerForTest(&RoundRobinSelector{}, httpOnly, wsA, wsB, wsC)
|
||||
|
||||
ctx := cliproxyexecutor.WithDownstreamWebsocket(context.Background())
|
||||
got, errPick := scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickSingle() first error = %v", errPick)
|
||||
}
|
||||
if got == nil || got.ID != "codex-ws-a" {
|
||||
t.Fatalf("pickSingle() first = %v, want codex-ws-a", got)
|
||||
}
|
||||
|
||||
wsA.Unavailable = true
|
||||
wsA.NextRetryAfter = time.Now().Add(time.Hour)
|
||||
scheduler.upsertAuth(wsA)
|
||||
|
||||
got, errPick = scheduler.pickSingle(ctx, "codex", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickSingle() after ws-a cooldown error = %v", errPick)
|
||||
}
|
||||
if got == nil || got.ID != "codex-ws-b" {
|
||||
t.Fatalf("pickSingle() after ws-a cooldown = %v, want codex-ws-b", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSchedulerPick_MixedProvidersPrefersHighestPriorityTier(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user