mirror of
https://gitcode.com/JianFeeeee/ModelRouter.git
synced 2026-09-20 00:48:00 +00:00
A slot that failed ten times landed in a 30-minute cooldown and was a hard
skip for the whole window: ModelAvailable() said no, the scheduler dropped the
candidate, and the ONLY way back was a success that could never happen. In
practice an upstream that blipped for a minute — or whose quota reset in
seconds — stayed unusable for half an hour.
Cooldown is now a window with a probe gate instead of a wall:
* backoffCap 30min -> 5min. 401/403 no longer implicitly "jump to the cap"
but get an explicit authCooldown (10min); 429 keeps its 30s fixed window.
* ModelState records cooldownFrom alongside cooldownUntil, so the window has
a measurable midpoint, plus a single-token probe permit (probeInFlight).
* TryProbe() hands out that permit only in the SECOND half of the window: a
freshly failed slot stays completely silent, and past the midpoint exactly
one request may go through as a probe.
* ProbeSlots() = clamp(max_concurrent/10, 1, 2), so a source configured for
max_concurrent=10 lets exactly one request probe.
* A probe is a REAL request: success runs RecordSuccess(), which clears the
cooldown and failCount on the spot and returns the slot to full rotation.
Failure reopens the window, pushing the next probe to the NEW midpoint
instead of retrying immediately.
Quota exhaustion is modelled separately (RecordQuotaExhausted): running out of
allowance is not the upstream being broken, so it must not feed the exponential
ladder or inflate failCount. It cools until the allowance window can plausibly
have reset (capped at 30min so a manual top-up is noticed) and is recognised
from the adapter-normalized error text via ReportStatusReason.
Scheduler side: collectCands() splits a tier into normal + probe candidates
with probes strictly LAST, so probe traffic is what a tier falls back to, never
what it prefers, and the round-robin cursor rotates over the normal head only.
releaseProbes() returns every permit on all exits (success, hard failure, ctx
cancel, tier fallthrough); the 90-line inlined busy-poll loop is extracted into
pollBusyTier() to keep those exits auditable. Direct Chat/ChatStream/Image take
the same gate.
Also fixes a real blacklist bug found on the way: a slot whose pref sank to
prefMin was refused by ModelAvailable forever, long after every cooldown had
expired. Such a slot is now probeable, and one success lifts it off the floor.
Verified live against a controllable upstream (tier1 flaky, tier2 healthy):
failure ladder 5s/9s/20s/40s/79s; after healing the upstream, 10 AUTO requests
across the silent half produced ZERO upstream calls; past the midpoint a single
AUTO request produced exactly ONE upstream call, cleared the remaining 39s of
cooldown, and the next three requests were served by the recovered slot.
498 lines
16 KiB
Go
498 lines
16 KiB
Go
package scheduler
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"llmsproxy/internal/types"
|
|
)
|
|
|
|
// fakeProvider is an in-memory Provider used to exercise chain scheduling
|
|
// deterministically without a Lua runtime.
|
|
type fakeProvider struct {
|
|
name string
|
|
model string
|
|
pref atomic.Int64
|
|
available atomic.Bool
|
|
busy atomic.Bool
|
|
fail atomic.Bool
|
|
chatHits atomic.Int64
|
|
// probeable makes an unavailable provider eligible as a half-cooldown
|
|
// probe candidate. probePermit is the single-token permit; probeClaims
|
|
// counts how many times it was handed out and probeDones how many times it
|
|
// was returned, so tests can assert the permit is always released.
|
|
probeable atomic.Bool
|
|
probePermit atomic.Bool
|
|
probeClaims atomic.Int64
|
|
probeDones atomic.Int64
|
|
}
|
|
|
|
func fakeProv(name, model string) *fakeProvider {
|
|
f := &fakeProvider{name: name, model: model}
|
|
f.available.Store(true)
|
|
return f
|
|
}
|
|
|
|
func (f *fakeProvider) Name() string { return f.name }
|
|
|
|
func (f *fakeProvider) ModelFor(reqModel string) string { return f.model }
|
|
|
|
func (f *fakeProvider) ModelAvailable(model string) bool { return f.available.Load() }
|
|
|
|
func (f *fakeProvider) ModelSchedulable(model string) (bool, bool) {
|
|
if f.available.Load() {
|
|
return true, false
|
|
}
|
|
if !f.probeable.Load() {
|
|
return false, false
|
|
}
|
|
if f.probePermit.CompareAndSwap(false, true) {
|
|
f.probeClaims.Add(1)
|
|
return true, true
|
|
}
|
|
return false, false
|
|
}
|
|
|
|
func (f *fakeProvider) ProbeDone(model string) {
|
|
f.probeDones.Add(1)
|
|
f.probePermit.Store(false)
|
|
}
|
|
|
|
func (f *fakeProvider) Pref(model string) int64 { return f.pref.Load() }
|
|
|
|
func (f *fakeProvider) Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error) {
|
|
f.chatHits.Add(1)
|
|
if f.busy.Load() {
|
|
return nil, types.ErrBusy
|
|
}
|
|
if f.fail.Load() {
|
|
return nil, fmt.Errorf("upstream error")
|
|
}
|
|
return &types.UnifiedResponse{Content: f.name, FinishReason: "stop", TokenUsage: types.TokenUsage{Prompt: 1, Completion: 1, Total: 2}}, nil
|
|
}
|
|
|
|
func (f *fakeProvider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error) {
|
|
if f.busy.Load() {
|
|
return nil, types.ErrBusy
|
|
}
|
|
if f.fail.Load() {
|
|
return nil, fmt.Errorf("upstream error")
|
|
}
|
|
ch := make(chan types.UnifiedChunk, 2)
|
|
ch <- types.UnifiedChunk{Content: f.name}
|
|
ch <- types.UnifiedChunk{Done: true}
|
|
close(ch)
|
|
return ch, nil
|
|
}
|
|
|
|
func (f *fakeProvider) Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error) {
|
|
return nil, errors.New("no image")
|
|
}
|
|
|
|
// TestDirectSkipsCooledCandidate: direct paths share the AUTO-chain rule that
|
|
// cooldown is the only hard skip — a cooling candidate must never be hit, and
|
|
// a fully cooling set must fail without touching upstream.
|
|
func TestDirectSkipsCooledCandidate(t *testing.T) {
|
|
hot := fakeProv("hot", "m")
|
|
cold := fakeProv("cold", "m")
|
|
cold.available.Store(false)
|
|
s := New(2)
|
|
req := &types.ChatRequest{
|
|
Model: "m",
|
|
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
|
|
}
|
|
resp, src, _, err := s.Chat(context.Background(), []Provider{cold, hot}, req)
|
|
if err != nil || src != "hot" {
|
|
t.Fatalf("want hot to serve, got src=%q err=%v", src, err)
|
|
}
|
|
if resp == nil || resp.Content != "hot" {
|
|
t.Fatalf("bad response: %#v", resp)
|
|
}
|
|
if cold.chatHits.Load() != 0 {
|
|
t.Fatalf("cooled candidate must not be hit, got %d hits", cold.chatHits.Load())
|
|
}
|
|
// all candidates cooled: direct path reports it instead of hitting upstream
|
|
hot.available.Store(false)
|
|
hits := hot.chatHits.Load()
|
|
_, _, _, err = s.Chat(context.Background(), []Provider{cold, hot}, req)
|
|
if err == nil || !strings.Contains(err.Error(), "cooling down") {
|
|
t.Fatalf("want cooling-down error, got %v", err)
|
|
}
|
|
if hot.chatHits.Load() != hits || cold.chatHits.Load() != 0 {
|
|
t.Fatalf("cooled candidates must not be hit, got hot=%d cold=%d", hot.chatHits.Load(), cold.chatHits.Load())
|
|
}
|
|
}
|
|
|
|
// lookup resolves (model, source) -> provider for chain builders in tests.
|
|
type lookup func(m, s string) Provider
|
|
|
|
func bySource(ps ...*fakeProvider) lookup {
|
|
return func(m, s string) Provider {
|
|
for _, p := range ps {
|
|
if p.name == s {
|
|
return p
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func chatReq() *types.ChatRequest {
|
|
return &types.ChatRequest{Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("hi")}}}
|
|
}
|
|
|
|
func TestBuildChainTierOrdering(t *testing.T) {
|
|
a, b, c, d := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c"), fakeProv("s4", "d")
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 1, Model: "a", Source: "s1"},
|
|
{Tier: 3, Model: "c", Source: "s3"},
|
|
{Tier: 2, Model: "b", Source: "s2"},
|
|
{Tier: 1, Model: "d", Source: "s4"},
|
|
{Tier: 9, Model: "gone", Source: "missing"},
|
|
}, bySource(a, b, c, d))
|
|
if len(ch.Tiers) != 3 {
|
|
t.Fatalf("tiers = %d, want 3", len(ch.Tiers))
|
|
}
|
|
if ch.Tiers[0].Tier != 1 || ch.Tiers[1].Tier != 2 || ch.Tiers[2].Tier != 3 {
|
|
t.Fatalf("tier order = %d,%d,%d, want 1,2,3", ch.Tiers[0].Tier, ch.Tiers[1].Tier, ch.Tiers[2].Tier)
|
|
}
|
|
// same-tier slots keep configured order, unresolvable rule is dropped
|
|
if len(ch.Tiers[0].Slots) != 2 || ch.Tiers[0].Slots[0].Model != "a" || ch.Tiers[0].Slots[1].Model != "d" {
|
|
t.Fatalf("tier 1 slots = %+v", ch.Tiers[0].Slots)
|
|
}
|
|
}
|
|
|
|
func TestChainRoundRobin(t *testing.T) {
|
|
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "a", Source: "s1"},
|
|
{Tier: 0, Model: "b", Source: "s2"},
|
|
}, bySource(a, b))
|
|
s := New(0)
|
|
var got []string
|
|
for i := 0; i < 4; i++ {
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("iter %d: %v", i, err)
|
|
}
|
|
got = append(got, src)
|
|
}
|
|
want := []string{"s1", "s2", "s1", "s2"}
|
|
for i := range want {
|
|
if got[i] != want[i] {
|
|
t.Fatalf("rr order = %v, want %v", got, want)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestChainPreferenceSinksButStaysReachable(t *testing.T) {
|
|
neg := fakeProv("neg", "n")
|
|
good := fakeProv("good", "g")
|
|
neg.pref.Store(-5)
|
|
good.fail.Store(true)
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "n", Source: "neg"},
|
|
{Tier: 0, Model: "g", Source: "good"},
|
|
}, bySource(neg, good))
|
|
s := New(0)
|
|
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("chain: %v", err)
|
|
}
|
|
// Round-robin starts at index 0 (neg). neg is available (pref=-5 > -20)
|
|
// and succeeds. good is never tried because neg already won.
|
|
if src != "neg" || model != "n" || resp.Content != "neg" {
|
|
t.Fatalf("served src=%q model=%q content=%q", src, model, resp.Content)
|
|
}
|
|
if good.chatHits.Load() != 0 {
|
|
t.Fatal("neg was tried first and succeeded; good must not be attempted")
|
|
}
|
|
// Second request: cursor advances. neg wins again (good hard-fails).
|
|
resp2, src2, _, err2 := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err2 != nil {
|
|
t.Fatalf("second chain: %v", err2)
|
|
}
|
|
if src2 != "neg" || resp2.Content != "neg" {
|
|
t.Fatalf("second req src=%q content=%q", src2, resp2.Content)
|
|
}
|
|
}
|
|
|
|
func TestChainBusySkipsWithoutPenalty(t *testing.T) {
|
|
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
|
|
a.busy.Store(true)
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "a", Source: "s1"},
|
|
{Tier: 0, Model: "b", Source: "s2"},
|
|
}, bySource(a, b))
|
|
s := New(0)
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("chain: %v", err)
|
|
}
|
|
if src != "s2" {
|
|
t.Fatalf("src = %q, want s2", src)
|
|
}
|
|
if a.chatHits.Load() == 0 {
|
|
t.Fatal("busy slot must have been attempted")
|
|
}
|
|
}
|
|
|
|
func TestChainAllBusyBoundedWaitThenNextTier(t *testing.T) {
|
|
oldWait, oldPoll := busyWait, busyPoll
|
|
busyWait, busyPoll = 60*time.Millisecond, 10*time.Millisecond
|
|
t.Cleanup(func() { busyWait, busyPoll = oldWait, oldPoll })
|
|
a, b, c := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c")
|
|
a.busy.Store(true)
|
|
b.busy.Store(true)
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 1, Model: "a", Source: "s1"},
|
|
{Tier: 1, Model: "b", Source: "s2"},
|
|
{Tier: 2, Model: "c", Source: "s3"},
|
|
}, bySource(a, b, c))
|
|
s := New(0)
|
|
t0 := time.Now()
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
el := time.Since(t0)
|
|
if err != nil {
|
|
t.Fatalf("chain: %v", err)
|
|
}
|
|
if src != "s3" {
|
|
t.Fatalf("src = %q, want s3 (downgrade after bounded wait)", src)
|
|
}
|
|
if el > time.Second {
|
|
t.Fatalf("busy wait not bounded: %v", el)
|
|
}
|
|
}
|
|
|
|
func TestChainQuotaExhausted(t *testing.T) {
|
|
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "a", Source: "s1", Quota: 100, Period: "hour"},
|
|
{Tier: 0, Model: "b", Source: "s2"},
|
|
}, bySource(a, b))
|
|
s := New(0)
|
|
exhausted := func(sl *Slot) bool { return sl.Source == "s1" && sl.Quota > 0 }
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), exhausted)
|
|
if err != nil {
|
|
t.Fatalf("chain: %v", err)
|
|
}
|
|
if src != "s2" {
|
|
t.Fatalf("src = %q, want s2 (quota slot dropped)", src)
|
|
}
|
|
if a.chatHits.Load() != 0 {
|
|
t.Fatal("quota-exhausted slot must not be called")
|
|
}
|
|
}
|
|
|
|
func TestChainErrSummary(t *testing.T) {
|
|
a, b, c := fakeProv("s1", "a"), fakeProv("s2", "b"), fakeProv("s3", "c")
|
|
a.fail.Store(true)
|
|
b.fail.Store(true)
|
|
c.available.Store(false) // whole tier 1 cooling
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "a", Source: "s1"},
|
|
{Tier: 0, Model: "b", Source: "s2"},
|
|
{Tier: 1, Model: "c", Source: "s3"},
|
|
}, bySource(a, b, c))
|
|
s := New(0)
|
|
_, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
var ce *ChainErr
|
|
if !errors.As(err, &ce) {
|
|
t.Fatalf("err = %v, want *ChainErr", err)
|
|
}
|
|
if len(ce.Tiers) != 2 || ce.Tiers[0].Source != "s1" || ce.Tiers[1].Source != "s2" {
|
|
t.Fatalf("tiers = %+v", ce.Tiers)
|
|
}
|
|
if len(ce.Skipped) != 1 {
|
|
t.Fatalf("skipped = %+v", ce.Skipped)
|
|
}
|
|
msg := ce.Error()
|
|
if !strings.Contains(msg, "s1") || !strings.Contains(msg, "s2") || !strings.Contains(msg, "tier 1: no schedulable slot") {
|
|
t.Fatalf("summary = %q", msg)
|
|
}
|
|
}
|
|
|
|
func TestChainStreamFallsBackBeforeFirstChunk(t *testing.T) {
|
|
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
|
|
a.fail.Store(true)
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "a", Source: "s1"},
|
|
{Tier: 0, Model: "b", Source: "s2"},
|
|
}, bySource(a, b))
|
|
s := New(0)
|
|
chunks, src, model, err := s.ChainChatStream(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("chain stream: %v", err)
|
|
}
|
|
if src != "s2" || model != "b" {
|
|
t.Fatalf("src=%q model=%q", src, model)
|
|
}
|
|
var text string
|
|
for ck := range chunks {
|
|
text += ck.Content
|
|
}
|
|
if text != "s2" {
|
|
t.Fatalf("text = %q", text)
|
|
}
|
|
}
|
|
|
|
func TestChainResetCooldownAfterSwap(t *testing.T) {
|
|
a := fakeProv("s1", "a")
|
|
ch := BuildChain([]Rule{{Tier: 0, Model: "a", Source: "s1"}}, bySource(a))
|
|
// a freshly built chain must schedule from index 0 (cursor starts at -1)
|
|
if base := ch.Tiers[0].NextStart(); base != 0 {
|
|
t.Fatalf("first start = %d, want 0", base)
|
|
}
|
|
}
|
|
|
|
// ---- half-cooldown probe candidates (plan 阶段 1.3) ----
|
|
|
|
// TestChainProbeIsLastResort pins the ordering rule: a probe candidate must not
|
|
// steal traffic from a healthy slot in the same tier.
|
|
func TestChainProbeIsLastResort(t *testing.T) {
|
|
healthy := fakeProv("healthy", "h")
|
|
cooling := fakeProv("cooling", "c")
|
|
cooling.available.Store(false)
|
|
cooling.probeable.Store(true)
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "c", Source: "cooling"}, // configured FIRST on purpose
|
|
{Tier: 0, Model: "h", Source: "healthy"},
|
|
}, bySource(healthy, cooling))
|
|
s := New(0)
|
|
for i := 0; i < 3; i++ {
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("iter %d: %v", i, err)
|
|
}
|
|
if src != "healthy" {
|
|
t.Fatalf("iter %d served by %q; a probe slot must never outrank a healthy one", i, src)
|
|
}
|
|
}
|
|
if cooling.chatHits.Load() != 0 {
|
|
t.Fatalf("probe slot was hit %d times while a healthy slot existed", cooling.chatHits.Load())
|
|
}
|
|
if got := cooling.probeDones.Load(); got != cooling.probeClaims.Load() {
|
|
t.Fatalf("probe permits leaked: claims=%d dones=%d", cooling.probeClaims.Load(), got)
|
|
}
|
|
}
|
|
|
|
// TestChainProbeServesWhenNothingElseCan is the recovery path: with every normal
|
|
// slot gone, the cooling slot's probe carries the request instead of the tier
|
|
// failing outright.
|
|
func TestChainProbeServesWhenNothingElseCan(t *testing.T) {
|
|
cooling := fakeProv("cooling", "c")
|
|
cooling.available.Store(false)
|
|
cooling.probeable.Store(true)
|
|
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
|
|
s := New(0)
|
|
resp, src, model, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("probe must serve the request: %v", err)
|
|
}
|
|
if src != "cooling" || model != "c" {
|
|
t.Fatalf("served by src=%q model=%q, want cooling/c", src, model)
|
|
}
|
|
if resp == nil || resp.Content != "cooling" {
|
|
t.Fatalf("bad response: %#v", resp)
|
|
}
|
|
if cooling.chatHits.Load() != 1 {
|
|
t.Fatalf("probe must issue exactly one request, got %d", cooling.chatHits.Load())
|
|
}
|
|
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
|
|
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
|
|
}
|
|
}
|
|
|
|
// TestChainProbePermitReleasedOnFailure: a failed probe must still hand its
|
|
// permit back, otherwise the slot could never be probed again.
|
|
func TestChainProbePermitReleasedOnFailure(t *testing.T) {
|
|
cooling := fakeProv("cooling", "c")
|
|
cooling.available.Store(false)
|
|
cooling.probeable.Store(true)
|
|
cooling.fail.Store(true)
|
|
ch := BuildChain([]Rule{{Tier: 0, Model: "c", Source: "cooling"}}, bySource(cooling))
|
|
s := New(0)
|
|
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
|
|
t.Fatal("expected the failing probe to surface an error")
|
|
}
|
|
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
|
|
t.Fatalf("permit accounting: claims=%d dones=%d, want 1/1", cooling.probeClaims.Load(), cooling.probeDones.Load())
|
|
}
|
|
// permit is free again for the next attempt
|
|
if _, _, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil); err == nil {
|
|
t.Fatal("expected the second probe to fail too")
|
|
}
|
|
if cooling.probeClaims.Load() != 2 {
|
|
t.Fatalf("permit must be reclaimable, claims=%d", cooling.probeClaims.Load())
|
|
}
|
|
}
|
|
|
|
// TestChainProbeDoesNotBlockTierFallthrough: a cooling tier-1 slot that is not
|
|
// probeable must still let the request fall through to tier 2.
|
|
func TestChainProbeDoesNotBlockTierFallthrough(t *testing.T) {
|
|
cold := fakeProv("cold", "c")
|
|
cold.available.Store(false) // cooling and NOT probeable
|
|
backup := fakeProv("backup", "b")
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 1, Model: "c", Source: "cold"},
|
|
{Tier: 2, Model: "b", Source: "backup"},
|
|
}, bySource(cold, backup))
|
|
s := New(0)
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil || src != "backup" {
|
|
t.Fatalf("want fallthrough to backup, got src=%q err=%v", src, err)
|
|
}
|
|
if cold.chatHits.Load() != 0 {
|
|
t.Fatal("non-probeable cooling slot must not be hit")
|
|
}
|
|
}
|
|
|
|
// TestDirectProbeReleasesPermit covers the non-AUTO paths (Chat/ChatStream/Image).
|
|
func TestDirectProbeReleasesPermit(t *testing.T) {
|
|
cooling := fakeProv("cooling", "m")
|
|
cooling.available.Store(false)
|
|
cooling.probeable.Store(true)
|
|
s := New(0)
|
|
_, src, _, err := s.Chat(context.Background(), []Provider{cooling}, chatReq())
|
|
if err != nil || src != "cooling" {
|
|
t.Fatalf("direct probe must serve: src=%q err=%v", src, err)
|
|
}
|
|
if cooling.probeClaims.Load() != 1 || cooling.probeDones.Load() != 1 {
|
|
t.Fatalf("permit accounting: claims=%d dones=%d", cooling.probeClaims.Load(), cooling.probeDones.Load())
|
|
}
|
|
}
|
|
|
|
// TestChainProbeRoundRobinUnaffected: adding a probe tail must not disturb the
|
|
// round-robin rotation over the healthy head.
|
|
func TestChainProbeRoundRobinUnaffected(t *testing.T) {
|
|
a, b := fakeProv("s1", "a"), fakeProv("s2", "b")
|
|
cooling := fakeProv("s3", "c")
|
|
cooling.available.Store(false)
|
|
cooling.probeable.Store(true)
|
|
ch := BuildChain([]Rule{
|
|
{Tier: 0, Model: "a", Source: "s1"},
|
|
{Tier: 0, Model: "b", Source: "s2"},
|
|
{Tier: 0, Model: "c", Source: "s3"},
|
|
}, bySource(a, b, cooling))
|
|
s := New(0)
|
|
var got []string
|
|
for i := 0; i < 4; i++ {
|
|
_, src, _, err := s.ChainChat(context.Background(), ch, chatReq(), nil)
|
|
if err != nil {
|
|
t.Fatalf("iter %d: %v", i, err)
|
|
}
|
|
got = append(got, src)
|
|
}
|
|
want := []string{"s1", "s2", "s1", "s2"}
|
|
for i := range want {
|
|
if got[i] != want[i] {
|
|
t.Fatalf("rr order = %v, want %v (probe tail must not join the rotation)", got, want)
|
|
}
|
|
}
|
|
}
|