fix(config): start model round robin from first match

This commit is contained in:
XYSK-lilong007 2026-03-12 08:28:32 +08:00
parent 4a8a2e9c23
commit 6dbc3bd645
2 changed files with 15 additions and 11 deletions

View file

@ -921,7 +921,7 @@ func (c *Config) GetModelConfig(modelName string) (*ModelConfig, error) {
} }
// Multiple configs - use round-robin for load balancing // Multiple configs - use round-robin for load balancing
idx := rrCounter.Add(1) % uint64(len(matches)) idx := (rrCounter.Add(1) - 1) % uint64(len(matches))
return &matches[idx], nil return &matches[idx], nil
} }

View file

@ -54,6 +54,8 @@ func TestGetModelConfig_EmptyList(t *testing.T) {
} }
func TestGetModelConfig_RoundRobin(t *testing.T) { func TestGetModelConfig_RoundRobin(t *testing.T) {
rrCounter.Store(0)
cfg := &Config{ cfg := &Config{
ModelList: []ModelConfig{ ModelList: []ModelConfig{
{ModelName: "lb-model", Model: "openai/gpt-4o-1", APIKey: "key1"}, {ModelName: "lb-model", Model: "openai/gpt-4o-1", APIKey: "key1"},
@ -62,20 +64,22 @@ func TestGetModelConfig_RoundRobin(t *testing.T) {
}, },
} }
// Test round-robin distribution want := []string{
results := make(map[string]int) "openai/gpt-4o-1",
for range 30 { "openai/gpt-4o-2",
"openai/gpt-4o-3",
"openai/gpt-4o-1",
"openai/gpt-4o-2",
"openai/gpt-4o-3",
}
for i, wantModel := range want {
result, err := cfg.GetModelConfig("lb-model") result, err := cfg.GetModelConfig("lb-model")
if err != nil { if err != nil {
t.Fatalf("GetModelConfig() error = %v", err) t.Fatalf("GetModelConfig() error = %v", err)
} }
results[result.Model]++ if result.Model != wantModel {
} t.Fatalf("call %d selected %q, want %q", i+1, result.Model, wantModel)
// Each model should appear roughly 10 times (30 calls / 3 models)
for model, count := range results {
if count < 5 || count > 15 {
t.Errorf("Model %s appeared %d times, expected ~10", model, count)
} }
} }
} }