mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-09-23 01:00:33 +08:00
feat(auth): enhance plugin scheduler with HasScheduler support and fast-path tests
- Added `pluginSchedulerState` interface with `HasScheduler` method for improved plugin scheduler state checks. - Updated `Manager.hasPluginScheduler` to handle `HasScheduler` logic. - Implemented and tested fast-path handling for inactive plugin schedulers, including mixed provider scenarios. - Expanded unit test coverage to ensure correct behavior in various scheduler states.
This commit is contained in:
@@ -30,6 +30,10 @@ func (h *Host) PickAuth(ctx context.Context, req pluginapi.SchedulerPickRequest)
|
||||
return resp, true, nil
|
||||
}
|
||||
|
||||
func (h *Host) HasScheduler() bool {
|
||||
return h.schedulerRecord() != nil
|
||||
}
|
||||
|
||||
func (h *Host) schedulerRecord() *capabilityRecord {
|
||||
if h == nil {
|
||||
return nil
|
||||
|
||||
@@ -126,6 +126,10 @@ type PluginScheduler interface {
|
||||
PickAuth(context.Context, pluginapi.SchedulerPickRequest) (pluginapi.SchedulerPickResponse, bool, error)
|
||||
}
|
||||
|
||||
type pluginSchedulerState interface {
|
||||
HasScheduler() bool
|
||||
}
|
||||
|
||||
// StoppableSelector is an optional interface for selectors that hold resources.
|
||||
// Selectors that implement this interface will have Stop called during shutdown.
|
||||
type StoppableSelector interface {
|
||||
@@ -245,9 +249,15 @@ func (m *Manager) hasPluginScheduler() bool {
|
||||
return false
|
||||
}
|
||||
m.mu.RLock()
|
||||
ok := m.pluginScheduler != nil
|
||||
scheduler := m.pluginScheduler
|
||||
m.mu.RUnlock()
|
||||
return ok
|
||||
if scheduler == nil {
|
||||
return false
|
||||
}
|
||||
if state, ok := scheduler.(pluginSchedulerState); ok {
|
||||
return state.HasScheduler()
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func isBuiltInSelector(selector Selector) bool {
|
||||
|
||||
@@ -55,6 +55,14 @@ func (s *fakePluginScheduler) PickAuth(ctx context.Context, req pluginapi.Schedu
|
||||
return s.resp, s.handled, s.err
|
||||
}
|
||||
|
||||
type inactivePluginScheduler struct {
|
||||
fakePluginScheduler
|
||||
}
|
||||
|
||||
func (s *inactivePluginScheduler) HasScheduler() bool {
|
||||
return false
|
||||
}
|
||||
|
||||
type trackingSelector struct {
|
||||
calls int
|
||||
lastAuthID []string
|
||||
@@ -440,6 +448,38 @@ func TestManagerPluginSchedulerSkippedWhenHomeEnabled(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerInactivePluginSchedulerKeepsFastPath(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-a) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "auth-b", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(auth-b) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &inactivePluginScheduler{}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
gotA, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() first error = %v", errPick)
|
||||
}
|
||||
gotB, _, errPick := manager.pickNext(context.Background(), "gemini", "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNext() second error = %v", errPick)
|
||||
}
|
||||
if gotA == nil || gotB == nil {
|
||||
t.Fatalf("pickNext() auths = %v, %v; want non-nil", gotA, gotB)
|
||||
}
|
||||
if gotA.ID != "auth-a" || gotB.ID != "auth-b" {
|
||||
t.Fatalf("fast path picks = %q, %q; want auth-a, auth-b", gotA.ID, gotB.ID)
|
||||
}
|
||||
if scheduler.calls != 0 {
|
||||
t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 0)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerCalledOutsideManagerLock(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
@@ -645,6 +685,38 @@ func TestManagerPluginSchedulerPickNextMixedSelectsProvider(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerInactivePluginSchedulerKeepsMixedFastPath(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
manager.executors["claude"] = schedulerTestExecutor{}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "gemini-a", Provider: "gemini"}); errRegister != nil {
|
||||
t.Fatalf("Register(gemini-a) error = %v", errRegister)
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), &Auth{ID: "claude-a", Provider: "claude"}); errRegister != nil {
|
||||
t.Fatalf("Register(claude-a) error = %v", errRegister)
|
||||
}
|
||||
|
||||
scheduler := &inactivePluginScheduler{}
|
||||
manager.SetPluginScheduler(scheduler)
|
||||
|
||||
got, _, provider, errPick := manager.pickNextMixed(context.Background(), []string{"gemini", "claude"}, "", cliproxyexecutor.Options{}, nil)
|
||||
if errPick != nil {
|
||||
t.Fatalf("pickNextMixed() error = %v", errPick)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("pickNextMixed() auth = nil")
|
||||
}
|
||||
if provider != "gemini" {
|
||||
t.Fatalf("pickNextMixed() provider = %q, want gemini", provider)
|
||||
}
|
||||
if got.ID != "gemini-a" {
|
||||
t.Fatalf("pickNextMixed() auth.ID = %q, want gemini-a", got.ID)
|
||||
}
|
||||
if scheduler.calls != 0 {
|
||||
t.Fatalf("scheduler.calls = %d, want %d", scheduler.calls, 0)
|
||||
}
|
||||
}
|
||||
|
||||
func TestManagerPluginSchedulerCandidatesAreSafeCopies(t *testing.T) {
|
||||
manager := NewManager(nil, &RoundRobinSelector{}, nil)
|
||||
manager.executors["gemini"] = schedulerTestExecutor{}
|
||||
|
||||
Reference in New Issue
Block a user