jianf/ModelRouter
1
0
Fork 0
mirror of https://gitcode.com/JianFeeeee/ModelRouter.git synced 2026-09-19 16:39:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(provider): half-cooldown probe so recovered slots return immediately

Browse Source 
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.
This commit is contained in:
JianFeeeee
2026-08-30 08:04:50 +08:00
parent 2ed1f0ecde
commit 21cd0429a2
4 changed files with 1011 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files

328
internal/provider/provider.go
View File

@ -13,6 +13,7 @@ import (
"log"
"net"
"net/http"
"strconv"
"strings"
"sync"
"sync/atomic"
@ -36,6 +37,16 @@ type ModelState struct {
pref atomic.Int64 // +1 per success / -5 per failure, clamped
failCount atomic.Int64
cooldownUntil atomic.Int64 // unix seconds; 0 = schedulable
// cooldownFrom is the unix second the current cooldown window opened. It
// is the reference point for the half-cooldown probe gate: probing is
// allowed only in the second half of [cooldownFrom, cooldownUntil), so a
// freshly failed slot stays fully silent for a while instead of being
// hammered immediately.
cooldownFrom atomic.Int64
// probeInFlight is the single-token probe permit. At most one probe
// request may be in flight per (source, model) while cooling; the token is
// released by ProbeDone regardless of outcome.
probeInFlight atomic.Bool
}
const (
@ -43,11 +54,33 @@ const (
prefMin = -20
prefMax = 20
backoffBase = 5 * time.Second
backoffCap = 30 * time.Minute
// backoffCap bounds the exponential schedule. It used to be 30min, which
// in practice meant "a slot whose upstream blipped is unusable for half an
// hour even though its quota recovered in seconds". Combined with the
// half-cooldown probe gate below a much shorter cap is safe: a still-dead
// upstream costs one probe request per cap/2, while a recovered one is back
// in full rotation on its first probe success.
backoffCap = 5 * time.Minute
// failures at/above this count back off at the capped duration
backoffCapN = 10
)
// authCooldown is the cooldown applied to 401/403 credential failures. A bad
// key does not fix itself in seconds, so it cools longer than a transport
// blip — but it is still a self-healing cooldown (a rotated key is picked up by
// the next probe), never a permanent blacklist.
const authCooldown = 10 * time.Minute
// maxQuotaCooldown bounds a quota-exhausted cooldown. Even when the reported
// reset window is huge (monthly quotas), the slot is re-probed at least this
// often so a manually topped-up allowance is noticed.
const maxQuotaCooldown = 30 * time.Minute
// quotaDefaultCooldown is used when the source gives no hint about its quota
// window length. It is long enough not to hammer an exhausted allowance and
// short enough that the half-cooldown probe retries within minutes.
const quotaDefaultCooldown = 10 * time.Minute
// rateLimitCooldown is the fixed cooldown applied on HTTP 429. A rate-limit
// rejection means "too fast", not "broken": unlike transport/5xx failures it
// must NOT escalate exponentially to backoffCap, or a quota-rich source gets
@ -81,25 +114,30 @@ func (s *ModelState) Available() bool {
}
// RecordFailure counts one consecutive failure and schedules exponential
// cooldown (5s, 10s, 20s … capped at 30min). auth marks 401/403 credential
// failures: it jumps straight to the capped cooldown and doubles the
// cooldown (5s, 10s, 20s … capped at backoffCap). auth marks 401/403
// credential failures: they get the longer fixed authCooldown and a doubled
// preference penalty, but the model still recovers when the cooldown expires.
func (s *ModelState) RecordFailure(auth bool) {
n := s.failCount.Add(1)
if auth && n < backoffCapN {
n = backoffCapN
s.failCount.Store(n) // persist the cap so FailCount() reports it too
}
var cd time.Duration
if n >= backoffCapN {
switch {
case auth:
// Persist the cap so FailCount() reports "this is a hard problem" and a
// following non-auth failure does not restart the ladder from the bottom.
if n < backoffCapN {
n = backoffCapN
s.failCount.Store(n)
}
cd = authCooldown
case n >= backoffCapN:
cd = backoffCap
} else {
default:
cd = backoffBase * time.Duration(1<<(n-1))
if cd > backoffCap {
cd = backoffCap
}
}
s.cooldownUntil.Store(time.Now().Add(cd).Unix())
s.openCooldown(cd)
pen := int64(prefFailStep)
if auth {
pen *= 2
@ -108,22 +146,57 @@ func (s *ModelState) RecordFailure(auth bool) {
clampPref(&s.pref, prefMin, prefMax)
}
// openCooldown starts a fresh cooldown window of length cd. Recording both
// ends of the window (not just the deadline) is what makes the half-cooldown
// probe gate possible, and resetting cooldownFrom on every new failure means a
// failed probe pushes the next probe to the middle of the NEW window instead of
// retrying immediately.
func (s *ModelState) openCooldown(cd time.Duration) {
now := time.Now()
s.cooldownFrom.Store(now.Unix())
s.cooldownUntil.Store(now.Add(cd).Unix())
}
// RecordRateLimit records an HTTP 429: a short fixed cooldown instead of the
// exponential schedule, so a merely-throttled source recovers in seconds and
// its remaining quota stays usable. The preference penalty still applies so
// other slots are preferred while cooling.
func (s *ModelState) RecordRateLimit() {
s.failCount.Add(1)
s.cooldownUntil.Store(time.Now().Add(rateLimitCooldown).Unix())
s.openCooldown(rateLimitCooldown)
s.pref.Add(-prefFailStep)
clampPref(&s.pref, prefMin, prefMax)
}
// RecordQuotaExhausted marks the model as out of allowance until its quota
// window resets. Unlike a failure this is not "the upstream is broken", so it
// must not feed the exponential ladder: the slot is unusable until the window
// boundary and immediately usable again after it. resetIn is the time left in
// the current quota window; non-positive values fall back to the short
// rate-limit window (the caller could not determine a period).
func (s *ModelState) RecordQuotaExhausted(resetIn time.Duration) {
if resetIn <= 0 {
resetIn = rateLimitCooldown
}
if resetIn > maxQuotaCooldown {
resetIn = maxQuotaCooldown
}
s.openCooldown(resetIn)
// No failCount bump: an exhausted quota is not an error streak, and
// letting it inflate failCount would make the NEXT real failure jump
// straight to the capped backoff.
s.pref.Add(-prefFailStep)
clampPref(&s.pref, prefMin, prefMax)
}
// RecordSuccess resets the failure counter and cooldown and bumps the
// preference score by one.
// preference score by one. A single success is enough to lift a slot off the
// prefMin floor (prefMin+1 > prefMin), which is what lets one successful probe
// return a written-off slot to normal rotation.
func (s *ModelState) RecordSuccess() {
s.failCount.Store(0)
s.cooldownUntil.Store(0)
s.cooldownFrom.Store(0)
s.pref.Add(1)
clampPref(&s.pref, prefMin, prefMax)
}
@ -131,9 +204,53 @@ func (s *ModelState) RecordSuccess() {
func (s *ModelState) reset() {
s.failCount.Store(0)
s.cooldownUntil.Store(0)
s.cooldownFrom.Store(0)
s.probeInFlight.Store(false)
s.pref.Store(0)
}
// TryProbe attempts to claim the single probe permit for a model that normal
// scheduling refuses. It returns true only when a probe is both needed and due:
//
// - cooling: allowed once the window passed its midpoint (the first half stays
// completely silent).
// - cooldown expired but the preference score sank to prefMin: ModelAvailable
// still refuses the slot, so a probe is the ONLY way back. Without this a
// slot that failed prefMin/prefFailStep times in a row would be blacklisted
// forever despite every cooldown having long expired.
//
// A healthy, non-cooling model returns false: the caller should schedule it
// normally. On true the caller MUST call ProbeDone exactly once.
func (s *ModelState) TryProbe() bool {
now := time.Now().Unix()
until := s.cooldownUntil.Load()
if until <= now {
if s.pref.Load() > prefMin {
return false // schedulable normally, no probe needed
}
// preference floor: probe to earn the way out
return s.probeInFlight.CompareAndSwap(false, true)
}
from := s.cooldownFrom.Load()
if from <= 0 || from >= until {
return false // no measurable window (legacy/degenerate state)
}
if now < from+(until-from)/2 {
return false // first half of the window: stay silent
}
return s.probeInFlight.CompareAndSwap(false, true)
}
// ProbeDone releases the probe permit claimed by TryProbe.
func (s *ModelState) ProbeDone() { s.probeInFlight.Store(false) }
// Probing reports whether a probe request currently holds the permit.
func (s *ModelState) Probing() bool { return s.probeInFlight.Load() }
// CooldownFrom is the unix timestamp the current cooldown window opened; 0
// when the model is not cooling.
func (s *ModelState) CooldownFrom() int64 { return s.cooldownFrom.Load() }
// Pref is the current preference score (higher = preferred).
func (s *ModelState) Pref() int64 { return s.pref.Load() }
@ -155,15 +272,15 @@ type Provider struct {
client *http.Client
stream *http.Client
mu sync.Mutex
sem chan struct{}
states map[string]*ModelState // key = model id
mu sync.Mutex
sem chan struct{}
states map[string]*ModelState // key = model id
// proactive rate limiting: nextOK is the earliest unix-nano time the next
// request may leave; throttle() spaces requests 60s/RPM apart. Zero when
// cfg.RPM <= 0 (unlimited).
rateMu sync.Mutex
nextOK int64
rpmGap time.Duration
rateMu sync.Mutex
nextOK int64
rpmGap time.Duration
lastProbe struct {
ok bool
err string
@ -370,6 +487,55 @@ func (p *Provider) ModelAvailable(model string) bool {
return true
}
// ModelSchedulable is the scheduling gate that replaces a bare ModelAvailable
// check. It reports whether the model may take a request right now and whether
// doing so is a cooldown PROBE rather than normal traffic:
//
// - (true, false): fully available, schedule normally.
// - (true, true): cooling, but past the half-cooldown mark and this caller
// claimed the probe permit. The caller MUST call ProbeDone(model) once the
// attempt finished, and SHOULD prefer non-probe candidates first.
// - (false, false): cooling and not probeable (first half of the window, the
// permit is taken, or the slot is persistently bad).
//
// Probing is what stops a recovered upstream from sitting out a whole cooldown:
// the probe is a real request, so success immediately clears the cooldown via
// RecordSuccess and the slot is back in full rotation.
func (p *Provider) ModelSchedulable(model string) (ok bool, isProbe bool) {
p.mu.Lock()
s, known := p.states[model]
p.mu.Unlock()
if !known {
return true, false
}
if s.Available() && s.Pref() > prefMin {
return true, false
}
if s.TryProbe() {
return true, true
}
return false, false
}
// ProbeDone releases a probe permit claimed through ModelSchedulable.
func (p *Provider) ProbeDone(model string) { p.state(model).ProbeDone() }
// ProbeSlots is how many concurrent probe requests this source tolerates while
// a model cools: one per 10 configured concurrency slots, at least 1 and never
// more than 2. A source configured for max_concurrent=10 therefore lets exactly
// one request through to probe, which is the intended "small-scale probe"
// behaviour; a large source still cannot flood a cooling upstream.
func (p *Provider) ProbeSlots() int {
n := p.cfg.MaxConcurrent / 10
if n < 1 {
return 1
}
if n > 2 {
return 2
}
return n
}
// Pref returns the adaptive preference score of a model (higher = preferred
// within a priority tier). Used by the AUTO chain to order same-tier slots.
func (p *Provider) Pref(model string) int64 {
@ -393,6 +559,19 @@ func (p *Provider) ModelHealthInfo(model string) (pref, failCount, cooldownUntil
return s.Pref(), s.FailCount(), s.CooldownUntil()
}
// ModelProbeInfo exposes the probe-gate state for the web UI: when the current
// cooldown window opened, when probing becomes allowed (its midpoint), and
// whether a probe is in flight right now. All zero when not cooling.
func (p *Provider) ModelProbeInfo(model string) (from, probeAfter int64, probing bool) {
s := p.state(model)
until := s.CooldownUntil()
from = s.CooldownFrom()
if until <= 0 || from <= 0 || from >= until {
return 0, 0, s.Probing()
}
return from, from + (until-from)/2, s.Probing()
}
// Probe performs a lightweight reachability + auth check against the source.
// It first tries GET <base>/models (fast, ~1s for OpenAI-compatible upstreams)
// and only falls back to a 1-token chat call when that endpoint is
@ -518,17 +697,33 @@ func (p *Provider) RecordSuccess(model string) {
}
// ReportStatus records an upstream HTTP status for the given model and drives
// the (source, model) backoff state. 401/403 → capped self-healing cooldown
// the (source, model) backoff state. 401/403 → long self-healing authCooldown
// with doubled penalty; 5xx → normal exponential backoff; 429 → short fixed
// cooldown (rateLimitCooldown): a rate rejection means "too fast", not
// "broken", so the slot must re-enter rotation quickly instead of escalating
// to a 30-minute lockout that wastes remaining quota.
// to a long lockout that wastes remaining quota.
//
// Prefer ReportStatusReason when the upstream error text is available: a
// quota-exhausted rejection is cooled to its window boundary instead of being
// treated as a failure streak.
func (p *Provider) ReportStatus(model string, code int) {
p.ReportStatusReason(model, code, "")
}
// ReportStatusReason is ReportStatus plus the upstream error text, which lets a
// quota/allowance rejection be distinguished from a generic rate limit. Quota
// exhaustion is not an error streak: it is cooled until the allowance window
// can plausibly have reset, and never inflates the exponential ladder.
func (p *Provider) ReportStatusReason(model string, code int, reason string) {
if code == 401 || code == 403 {
p.RecordFailure(model, code)
return
}
if code == 429 {
if code == 429 || code == 402 {
if quotaExhaustedReason(reason) {
p.state(model).RecordQuotaExhausted(p.quotaResetIn(model))
return
}
p.state(model).RecordRateLimit()
return
}
@ -537,6 +732,82 @@ func (p *Provider) ReportStatus(model string, code int) {
}
}
// quotaResetIn estimates how long the model's allowance window still has to
// run. Upstreams do not report this, so it is taken from the source's optional
// `quota_hours` meta hint when present (aligned to the epoch, matching the
// gateway's own quota-window convention) and otherwise falls back to
// quotaDefaultCooldown. Either way the value is capped by maxQuotaCooldown, so
// a wrong hint can only delay recovery to that bound — and the half-cooldown
// probe fires at half of it.
func (p *Provider) quotaResetIn(model string) time.Duration {
hours := p.quotaHours()
if hours <= 0 {
return quotaDefaultCooldown
}
win := int64(hours * float64(time.Hour/time.Second))
if win <= 0 {
return quotaDefaultCooldown
}
elapsed := time.Now().Unix() % win
return time.Duration(win-elapsed) * time.Second
}
// quotaHours reads the optional per-source `quota_hours` meta hint describing
// how long its allowance window is. Absent or unparseable = 0 (unknown).
func (p *Provider) quotaHours() float64 {
v, ok := p.cfg.Meta["quota_hours"]
if !ok {
return 0
}
switch n := v.(type) {
case float64:
return n
case int:
return float64(n)
case int64:
return float64(n)
case string:
f, err := strconv.ParseFloat(strings.TrimSpace(n), 64)
if err != nil {
return 0
}
return f
}
return 0
}
// quotaKeywords are the upstream phrases that mean "you are out of allowance"
// rather than "you are going too fast". Matching is case-insensitive substring
// matching against the adapter-normalized error reason.
var quotaKeywords = []string{
"insufficient_quota",
"insufficient quota",
"quota exceeded",
"quota_exceeded",
"exceeded your current quota",
"out of credit",
"insufficient balance",
"insufficient_user_quota",
"billing_hard_limit_reached",
"credit balance is too low",
"余额不足",
"额度不足",
"额度已用完",
}
func quotaExhaustedReason(reason string) bool {
if reason == "" {
return false
}
low := strings.ToLower(reason)
for _, k := range quotaKeywords {
if strings.Contains(low, k) {
return true
}
}
return false
}
// ResetHealth resets the scheduling state of every model of this source
// (cooldown and preference to zero), so the source becomes fully schedulable
// again. Called after AUTO-chain edits and from the admin UI.
@ -717,8 +988,9 @@ func (p *Provider) Chat(ctx context.Context, req *types.ChatRequest) (*types.Uni
return nil, err
}
if status != 200 {
p.ReportStatus(model, status)
return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
reason := p.apiErrReason(status, raw)
p.ReportStatusReason(model, status, reason)
return nil, fmt.Errorf("%s", reason)
}
unified, err := p.vm.Transform(p.adapter, "transform_response", raw)
if err != nil {
@ -792,9 +1064,10 @@ func (p *Provider) ChatStream(ctx context.Context, req *types.ChatRequest) (<-ch
if sel.resp.StatusCode != 200 {
raw, _ := io.ReadAll(sel.resp.Body)
sel.resp.Body.Close()
p.ReportStatus(model, sel.resp.StatusCode)
reason := p.apiErrReason(sel.resp.StatusCode, string(raw))
p.ReportStatusReason(model, sel.resp.StatusCode, reason)
p.Release()
return nil, fmt.Errorf("%s", p.apiErrReason(sel.resp.StatusCode, string(raw)))
return nil, fmt.Errorf("%s", reason)
}
go func() {
defer p.Release()
@ -1043,8 +1316,9 @@ func (p *Provider) Image(ctx context.Context, req *types.ImageGenRequest) (*type
return nil, err
}
if status != 200 {
p.ReportStatus(model, status)
return nil, fmt.Errorf("%s", p.apiErrReason(status, raw))
reason := p.apiErrReason(status, raw)
p.ReportStatusReason(model, status, reason)
return nil, fmt.Errorf("%s", reason)
}
var out types.UnifiedResponse
// try adapter transform_response; if missing, parse standard openai image format

370
internal/provider/provider_test.go
View File

@ -198,8 +198,8 @@ func TestProviderAuthFailureSelfHeals(t *testing.T) {
if st.FailCount() != backoffCapN {
t.Fatalf("auth failure must jump to capped count, got %d", st.FailCount())
}
if until := st.CooldownUntil(); until <= time.Now().Add(25*time.Minute).Unix() {
t.Fatalf("auth failure must cool near the cap (until=%d)", until)
if until := st.CooldownUntil(); until <= time.Now().Add(authCooldown-time.Minute).Unix() {
t.Fatalf("auth failure must cool for authCooldown (until=%d)", until)
}
if st.Pref() != -2*int64(prefFailStep) {
t.Fatalf("auth failure pref penalty must be doubled, got %d", st.Pref())
@ -678,3 +678,369 @@ func TestToolIndexRemapper(t *testing.T) {
}
})
}
// ---- half-cooldown probe gate (plan 阶段 1) ----
// TestProbeGateHalfCooldown pins the core rule: a cooling model is fully silent
// in the first half of its cooldown window and probeable exactly once in the
// second half.
func TestProbeGateHalfCooldown(t *testing.T) {
var st ModelState
// open a 100s window centred so that "now" sits in the first half
now := time.Now().Unix()
st.cooldownFrom.Store(now - 10)
st.cooldownUntil.Store(now + 90)
if st.TryProbe() {
t.Fatal("first half of the cooldown window must stay silent")
}
// move into the second half
st.cooldownFrom.Store(now - 90)
st.cooldownUntil.Store(now + 10)
if !st.TryProbe() {
t.Fatal("second half of the cooldown window must allow one probe")
}
if st.TryProbe() {
t.Fatal("probe permit must be single-token")
}
if !st.Probing() {
t.Fatal("Probing() must report the claimed permit")
}
st.ProbeDone()
if st.Probing() {
t.Fatal("ProbeDone must release the permit")
}
if !st.TryProbe() {
t.Fatal("permit must be reusable after ProbeDone")
}
}
// TestProbeGateNoProbeWhenHealthy: a healthy model needs no probe permit — the
// caller schedules it normally.
func TestProbeGateNoProbeWhenHealthy(t *testing.T) {
var idle ModelState
if idle.TryProbe() {
t.Fatal("a healthy model must not be probed (schedule it normally)")
}
}
// TestProbeGateRescuesPrefFloor documents a fix that comes with the probe gate:
// a slot whose preference sank to prefMin used to be refused by ModelAvailable
// forever, even long after every cooldown expired — a permanent blacklist. It is
// now probeable, so one success brings it back.
func TestProbeGateRescuesPrefFloor(t *testing.T) {
var dead ModelState
dead.pref.Store(prefMin)
if !dead.TryProbe() {
t.Fatal("a slot stuck at prefMin must be probeable, not blacklisted")
}
dead.RecordSuccess()
dead.ProbeDone()
if dead.Pref() <= prefMin {
t.Fatalf("one probe success must lift the slot off the floor, pref=%d", dead.Pref())
}
}
// TestProbeSuccessRestoresImmediately is the user-visible payoff: a recovered
// upstream must not sit out the rest of its cooldown.
func TestProbeSuccessRestoresImmediately(t *testing.T) {
var st ModelState
for i := 0; i < backoffCapN; i++ {
st.RecordFailure(false)
}
if cd := st.CooldownUntil() - time.Now().Unix(); cd < int64(backoffCap.Seconds())-2 {
t.Fatalf("10 failures must cool for backoffCap, got %ds", cd)
}
// jump into the probe window: the midpoint of [from, until) must be in the
// past, so push the window's start back past its full length
now := time.Now().Unix()
st.cooldownFrom.Store(now - int64(backoffCap.Seconds()) - 10)
if !st.TryProbe() {
t.Fatal("probe must be allowed past the midpoint")
}
if !st.Probing() {
t.Fatal("probe permit must be held while the probe runs")
}
st.RecordSuccess()
st.ProbeDone()
if !st.Available() {
t.Fatal("probe success must clear the cooldown immediately")
}
if st.FailCount() != 0 {
t.Fatalf("probe success must reset failCount, got %d", st.FailCount())
}
if st.CooldownUntil() != 0 || st.CooldownFrom() != 0 {
t.Fatal("probe success must clear both cooldown bounds")
}
}
// TestProbeFailureDefersNextProbe: a failed probe must not immediately retry —
// the new window's midpoint pushes the next probe out.
func TestProbeFailureDefersNextProbe(t *testing.T) {
var st ModelState
st.RecordFailure(false) // 5s window
now := time.Now().Unix()
st.cooldownFrom.Store(now - 100)
st.cooldownUntil.Store(now + 10)
if !st.TryProbe() {
t.Fatal("probe must be allowed past the midpoint")
}
st.RecordFailure(false) // probe failed: a fresh, longer window opens
st.ProbeDone()
if st.TryProbe() {
t.Fatal("a failed probe must defer the next probe to the new midpoint")
}
if from := st.CooldownFrom(); from < now {
t.Fatalf("failed probe must reopen the window (from=%d, now=%d)", from, now)
}
}
// TestAuthCooldownIsSeparateFromCap documents that credential failures cool on
// their own (longer) schedule instead of reusing the exponential cap.
func TestAuthCooldownIsSeparateFromCap(t *testing.T) {
var st ModelState
st.RecordFailure(true)
cd := st.CooldownUntil() - time.Now().Unix()
if cd < int64(authCooldown.Seconds())-2 || cd > int64(authCooldown.Seconds())+2 {
t.Fatalf("auth failure must cool for authCooldown (%v), got %ds", authCooldown, cd)
}
if st.FailCount() != backoffCapN {
t.Fatalf("auth failure must persist the capped count, got %d", st.FailCount())
}
}
// TestQuotaExhaustedDoesNotEscalate is the second half of the user's complaint:
// running out of allowance must not be treated as a failure streak.
func TestQuotaExhaustedDoesNotEscalate(t *testing.T) {
var st ModelState
st.RecordQuotaExhausted(2 * time.Minute)
if st.FailCount() != 0 {
t.Fatalf("quota exhaustion must not bump failCount, got %d", st.FailCount())
}
cd := st.CooldownUntil() - time.Now().Unix()
if cd < 118 || cd > 122 {
t.Fatalf("quota cooldown must honour the reset window, got %ds", cd)
}
// an unbounded window is clamped so a topped-up allowance is still noticed
var big ModelState
big.RecordQuotaExhausted(72 * time.Hour)
if cd := big.CooldownUntil() - time.Now().Unix(); cd > int64(maxQuotaCooldown.Seconds())+2 {
t.Fatalf("quota cooldown must be capped at maxQuotaCooldown, got %ds", cd)
}
// no hint at all falls back to the short window
var none ModelState
none.RecordQuotaExhausted(0)
if cd := none.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
t.Fatalf("unknown reset window must fall back to rateLimitCooldown, got %ds", cd)
}
}
// TestQuotaReasonRouting checks that a 429 carrying a quota phrase is cooled as
// quota exhaustion while a plain rate limit keeps the short fixed window.
func TestQuotaReasonRouting(t *testing.T) {
p := newTestProvider(t, src("q", "http://127.0.0.1:1", "openai", "m"))
p.ReportStatusReason("m", 429, "429 insufficient_quota: you exceeded your current quota")
st := p.state("m")
if st.FailCount() != 0 {
t.Fatalf("quota rejection must not bump failCount, got %d", st.FailCount())
}
if cd := st.CooldownUntil() - time.Now().Unix(); cd <= int64(rateLimitCooldown.Seconds()) {
t.Fatalf("quota rejection must cool longer than a plain 429, got %ds", cd)
}
p2 := newTestProvider(t, src("r", "http://127.0.0.1:1", "openai", "m"))
p2.ReportStatusReason("m", 429, "rate limit reached, slow down")
st2 := p2.state("m")
if cd := st2.CooldownUntil() - time.Now().Unix(); cd > int64(rateLimitCooldown.Seconds())+2 {
t.Fatalf("plain 429 must keep the short window, got %ds", cd)
}
}
// TestModelSchedulableProbeHandoff covers the provider-level gate the scheduler
// consumes.
func TestModelSchedulableProbeHandoff(t *testing.T) {
p := newTestProvider(t, src("s", "http://127.0.0.1:1", "openai", "m"))
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
t.Fatalf("healthy model must be normally schedulable, got ok=%v probe=%v", ok, probe)
}
st := p.state("m")
st.RecordFailure(false)
now := time.Now().Unix()
st.cooldownFrom.Store(now - 100)
st.cooldownUntil.Store(now + 10)
ok, probe := p.ModelSchedulable("m")
if !ok || !probe {
t.Fatalf("cooling model past midpoint must be probeable, got ok=%v probe=%v", ok, probe)
}
if ok2, _ := p.ModelSchedulable("m"); ok2 {
t.Fatal("only one probe permit may be outstanding")
}
p.ProbeDone("m")
if ok3, probe3 := p.ModelSchedulable("m"); !ok3 || !probe3 {
t.Fatal("permit must be reclaimable after ProbeDone")
}
}
// TestProbeSlotsFollowsMaxConcurrent pins the user's example: max_concurrent=10
// yields exactly one probe slot.
func TestProbeSlotsFollowsMaxConcurrent(t *testing.T) {
for _, tc := range []struct{ mc, want int }{
{0, 1}, {1, 1}, {4, 1}, {10, 1}, {20, 2}, {100, 2},
} {
s := src("p", "http://127.0.0.1:1", "openai", "m")
s.MaxConcurrent = tc.mc
p := newTestProvider(t, s)
if got := p.ProbeSlots(); got != tc.want {
t.Fatalf("max_concurrent=%d: probe slots = %d, want %d", tc.mc, got, tc.want)
}
}
}
// TestModelProbeInfo checks the UI-facing view of the probe window.
func TestModelProbeInfo(t *testing.T) {
p := newTestProvider(t, src("i", "http://127.0.0.1:1", "openai", "m"))
if from, after, probing := p.ModelProbeInfo("m"); from != 0 || after != 0 || probing {
t.Fatalf("idle model must report no probe window, got %d %d %v", from, after, probing)
}
st := p.state("m")
now := time.Now().Unix()
st.cooldownFrom.Store(now)
st.cooldownUntil.Store(now + 100)
from, after, _ := p.ModelProbeInfo("m")
if from != now || after != now+50 {
t.Fatalf("probe window = (%d,%d), want (%d,%d)", from, after, now, now+50)
}
}
// TestUpstreamRecoveryWithoutWaitingOutCooldown is the end-to-end shape of the
// user's complaint: an upstream that 500s ten times lands in the capped
// cooldown, and used to be unusable for a full 30 minutes even after it
// recovered. With the probe gate the slot is silent for the first half of the
// window and then serves the very next request itself, clearing the cooldown.
//
// Real time is not waited out: the window bounds are moved so the test observes
// the same state the scheduler would see at cap/2.
func TestUpstreamRecoveryWithoutWaitingOutCooldown(t *testing.T) {
var healthy atomic.Bool
var hits atomic.Int64
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
if !healthy.Load() {
w.WriteHeader(500)
return
}
fmt.Fprint(w, `{"choices":[{"message":{"content":"back"},"finish_reason":"stop"}]}`)
}))
defer up.Close()
p := newTestProvider(t, src("recov", up.URL, "openai", "m"))
req := func() *types.ChatRequest {
return &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
}
}
// 1. drive the slot into the capped cooldown
for i := 0; i < backoffCapN; i++ {
if _, err := p.Chat(context.Background(), req()); err == nil {
t.Fatalf("attempt %d: expected upstream failure", i)
}
}
st := p.state("m")
if p.ModelAvailable("m") {
t.Fatal("10 consecutive 5xx must take the slot out of normal rotation")
}
cd := st.CooldownUntil() - time.Now().Unix()
if cd < int64(backoffCap.Seconds())-2 {
t.Fatalf("cooldown = %ds, want ~%v", cd, backoffCap)
}
// the old 30-minute lockout is gone
if cd > int64((6 * time.Minute).Seconds()) {
t.Fatalf("cooldown %ds is too long: the cap must be minutes, not half an hour", cd)
}
// 2. upstream recovers, but we are still in the FIRST half of the window:
// the slot must stay completely silent (no probe traffic at all).
healthy.Store(true)
quiet := hits.Load()
if ok, probe := p.ModelSchedulable("m"); ok || probe {
t.Fatalf("first half of the window must stay silent, got ok=%v probe=%v", ok, probe)
}
if hits.Load() != quiet {
t.Fatal("no request may reach the upstream during the silent half")
}
// 3. advance past the midpoint: exactly one probe is admitted.
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
ok, isProbe := p.ModelSchedulable("m")
if !ok || !isProbe {
t.Fatalf("past the midpoint one probe must be admitted, got ok=%v probe=%v", ok, isProbe)
}
if ok2, _ := p.ModelSchedulable("m"); ok2 {
t.Fatal("only ONE probe may be in flight while cooling")
}
// 4. the probe is a real request: its success restores the slot fully.
resp, err := p.Chat(context.Background(), req())
p.ProbeDone("m")
if err != nil {
t.Fatalf("probe request failed: %v", err)
}
if resp.Content != "back" {
t.Fatalf("probe response = %q, want \"back\"", resp.Content)
}
if !p.ModelAvailable("m") {
t.Fatal("a successful probe must return the slot to normal rotation immediately")
}
if st.CooldownUntil() != 0 || st.FailCount() != 0 {
t.Fatalf("probe success must clear cooldown/failCount, got until=%d fails=%d",
st.CooldownUntil(), st.FailCount())
}
if ok, probe := p.ModelSchedulable("m"); !ok || probe {
t.Fatalf("recovered slot must schedule normally, got ok=%v probe=%v", ok, probe)
}
}
// TestRecoveryProbeBudgetIsBounded quantifies the cost of the shorter cap
// against a still-dead upstream: one probe per window, not a retry storm.
func TestRecoveryProbeBudgetIsBounded(t *testing.T) {
var hits atomic.Int64
up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
hits.Add(1)
w.WriteHeader(500)
}))
defer up.Close()
p := newTestProvider(t, src("dead", up.URL, "openai", "m"))
req := &types.ChatRequest{
Model: "m",
Messages: []types.ChatMessage{{Role: "user", Content: types.StringContent("x")}},
}
for i := 0; i < backoffCapN; i++ {
_, _ = p.Chat(context.Background(), req)
}
baseline := hits.Load()
st := p.state("m")
// simulate three consecutive probe windows against a dead upstream
for i := 0; i < 3; i++ {
st.cooldownFrom.Store(time.Now().Unix() - int64(backoffCap.Seconds()) - 1)
ok, isProbe := p.ModelSchedulable("m")
if !ok || !isProbe {
t.Fatalf("window %d: expected a probe permit", i)
}
// a second caller in the same window gets nothing
if ok2, _ := p.ModelSchedulable("m"); ok2 {
t.Fatalf("window %d: more than one probe admitted", i)
}
if _, err := p.Chat(context.Background(), req); err == nil {
t.Fatalf("window %d: expected the probe to fail", i)
}
p.ProbeDone("m")
}
if extra := hits.Load() - baseline; extra != 3 {
t.Fatalf("3 probe windows must cost exactly 3 upstream requests, got %d", extra)
}
}

233
internal/scheduler/scheduler.go
View File

@ -42,6 +42,12 @@ type Provider interface {
Name() string
ModelFor(reqModel string) string
ModelAvailable(model string) bool
// ModelSchedulable is the probe-aware availability gate: ok reports
// whether the model may take a request now, isProbe marks that it is only
// allowed as a cooldown probe (the caller must release the permit with
// ProbeDone once the attempt finished).
ModelSchedulable(model string) (ok bool, isProbe bool)
ProbeDone(model string)
Pref(model string) int64
Chat(ctx context.Context, req *types.ChatRequest) (*types.UnifiedResponse, error)
ChatStream(ctx context.Context, req *types.ChatRequest) (<-chan types.UnifiedChunk, error)
@ -176,18 +182,84 @@ type tierResult struct {
hard []TierError // hard failures seen in this pass (nil = none)
}
// candidate is one schedulable slot in a tier pass. probe marks a slot that is
// still cooling but past its half-cooldown mark and holding the probe permit:
// it is tried only after every normal slot, so probe traffic is what the tier
// falls back to instead of what it prefers.
type candidate struct {
slot *Slot
probe bool
}
// collectCands partitions a tier's slots into normal and probe candidates,
// normal first. Quota-exhausted slots are dropped outright. Every probe
// candidate returned holds a probe permit, so the caller MUST call
// releaseProbes on the result exactly once.
func collectCands(slots []*Slot, exhausted func(*Slot) bool) []candidate {
var normal, probes []candidate
for _, sl := range slots {
if exhausted != nil && exhausted(sl) {
continue
}
ok, isProbe := sl.Prov.ModelSchedulable(sl.Model)
if !ok {
continue
}
if isProbe {
probes = append(probes, candidate{slot: sl, probe: true})
continue
}
normal = append(normal, candidate{slot: sl})
}
return append(normal, probes...)
}
// releaseProbes hands every claimed probe permit back, whether or not the probe
// slot was actually used.
func releaseProbes(cands []candidate) {
for _, c := range cands {
if c.probe {
c.slot.Prov.ProbeDone(c.slot.Model)
}
}
}
// normalCount is how many leading candidates are normal (non-probe). The
// round-robin cursor rotates only over those: probe slots are a strictly
// ordered tail, never a rotation target.
func normalCount(cands []candidate) int {
for i, c := range cands {
if c.probe {
return i
}
}
return len(cands)
}
// runTier executes one tier pass starting at the round-robin base index.
// Cooldown is the only hard skip (re-verified per slot); a busy slot is
// skipped without any penalty; a hard failure is recorded and the pass moves
// on to the next slot (plan 2.3: "单请求内不重试已失败槽" — the failed slot is
// not retried, the others still are). hard == nil and no success means every
// candidate was merely busy/cooling, so the caller may wait a bounded time.
func runTier(ctx context.Context, tn *TierNode, cands []*Slot, base int64, req *types.ChatRequest, stream bool) tierResult {
//
// Normal candidates rotate by base; probe candidates form a fixed tail tried
// only after every normal slot failed or was busy.
func runTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool) tierResult {
n := len(cands)
norm := normalCount(cands)
var hard []TierError
for i := 0; i < n; i++ {
sl := cands[(int(base)+i)%n]
if !sl.Prov.ModelAvailable(sl.Model) {
var c candidate
if i < norm {
c = cands[(int(base)+i)%norm] // rotate within the normal head
} else {
c = cands[i] // probe tail keeps its order
}
sl := c.slot
// A probe candidate is intentionally NOT re-checked here: it is cooling
// by definition, and its permit was already claimed.
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
@ -233,75 +305,40 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
}
var ce ChainErr
for _, tn := range chain.Tiers {
// initial filter: quota-exhausted and cooling slots are dropped
var cands []*Slot
for _, sl := range tn.Slots {
if exhausted != nil && exhausted(sl) {
continue
}
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
cands = append(cands, sl)
}
// initial filter: quota-exhausted slots are dropped, cooling slots are
// dropped unless they qualify as half-cooldown probes (appended last).
cands := collectCands(tn.Slots, exhausted)
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling or quota exhausted)", tn.Tier))
continue
}
// No Pref sort: load balancing is done by round-robin cursor.
// Persistently failing slots are excluded by ModelAvailable
// Persistently failing slots are excluded by ModelSchedulable
// (which checks Pref > prefMin).
base := tn.NextStart()
res := runTier(ctx, tn, cands, base, req, stream)
if res.resp != nil || res.chunks != nil {
releaseProbes(cands)
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
releaseProbes(cands)
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
releaseProbes(cands)
continue // hard failures: fall through to the next tier, no waiting
}
// every candidate was busy or cooling: bounded poll before downgrading
deadline := time.Now().Add(busyWait)
timer := time.NewTimer(busyPoll)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return nil, nil, "", "", ctx.Err()
case <-timer.C:
if err := s.pollBusyTier(ctx, tn, cands, base, req, stream, &ce); err != nil {
releaseProbes(cands)
if r, ok := err.(*tierSuccess); ok {
return r.res.resp, r.res.chunks, r.res.src, r.res.model, nil
}
done := time.Now().After(deadline)
if done {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
// refresh candidates: cooldowns may have expired meanwhile
var again []*Slot
for _, sl := range cands {
if sl.Prov.ModelAvailable(sl.Model) {
again = append(again, sl)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
break
}
res = runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return res.resp, res.chunks, res.src, res.model, nil
}
if ctx.Err() != nil {
return nil, nil, "", "", ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
break // hard failure while waiting: stop waiting, fall through
}
timer.Reset(busyPoll)
return nil, nil, "", "", err
}
releaseProbes(cands)
}
if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
return nil, nil, "", "", fmt.Errorf("no auto slot configured")
@ -309,6 +346,57 @@ func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.Cha
return nil, nil, "", "", &ce
}
// tierSuccess carries a successful result out of pollBusyTier through the error
// return. It is never surfaced to callers of chainDrive.
type tierSuccess struct{ res tierResult }
func (t *tierSuccess) Error() string { return "tier success" }
// pollBusyTier waits a bounded time for a fully-busy tier to free a slot,
// retrying the pass while cooldowns expire. It returns nil when the tier should
// be abandoned (caller falls through to the next tier), a *tierSuccess when a
// retry succeeded, or a context error.
func (s *Scheduler) pollBusyTier(ctx context.Context, tn *TierNode, cands []candidate, base int64, req *types.ChatRequest, stream bool, ce *ChainErr) error {
deadline := time.Now().Add(busyWait)
timer := time.NewTimer(busyPoll)
defer timer.Stop()
for {
select {
case <-ctx.Done():
return ctx.Err()
case <-timer.C:
}
if time.Now().After(deadline) {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
return nil
}
// refresh candidates: cooldowns may have expired meanwhile. Probe
// candidates keep their already-claimed permit and stay eligible.
var again []candidate
for _, c := range cands {
if c.probe || c.slot.Prov.ModelAvailable(c.slot.Model) {
again = append(again, c)
}
}
if len(again) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no free slot within %v", tn.Tier, busyWait))
return nil
}
res := runTier(ctx, tn, again, base, req, stream)
if res.resp != nil || res.chunks != nil {
return &tierSuccess{res: res}
}
if ctx.Err() != nil {
return ctx.Err()
}
if len(res.hard) > 0 {
ce.Tiers = append(ce.Tiers, res.hard...)
return nil // hard failure while waiting: stop waiting, fall through
}
timer.Reset(busyPoll)
}
}
// ChainChat runs a non-streaming AUTO request down the chain. exhausted, when
// non-nil, decides slot token-quota exhaustion. Returns the response, the
// serving source and the exact model id used; on total failure a *ChainErr
@ -338,32 +426,35 @@ func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.Ima
}
var ce ChainErr
for _, tn := range chain.Tiers {
var cands []*Slot
for _, sl := range tn.Slots {
if !sl.Prov.ModelAvailable(sl.Model) {
continue
}
cands = append(cands, sl)
}
cands := collectCands(tn.Slots, nil)
if len(cands) == 0 {
ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling)", tn.Tier))
continue
}
// No Pref sort: load balancing is done by round-robin cursor.
base := tn.NextStart()
norm := normalCount(cands)
var hard []TierError
for i := 0; i < len(cands); i++ {
sl := cands[(int(base)+i)%len(cands)]
if !sl.Prov.ModelAvailable(sl.Model) {
var c candidate
if i < norm {
c = cands[(int(base)+i)%norm]
} else {
c = cands[i]
}
sl := c.slot
if !c.probe && !sl.Prov.ModelAvailable(sl.Model) {
continue
}
r := *req
r.Model = sl.Model
resp, err := sl.Prov.Image(ctx, &r)
if ctx.Err() != nil {
releaseProbes(cands)
return nil, "", "", ctx.Err()
}
if err == nil {
releaseProbes(cands)
return resp, sl.Source, sl.Model, nil
}
if errors.Is(err, types.ErrBusy) {
@ -371,6 +462,7 @@ func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.Ima
}
hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
}
releaseProbes(cands)
if len(hard) > 0 {
ce.Tiers = append(ce.Tiers, hard...)
}
@ -399,11 +491,15 @@ func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatR
p := cands[i]
r := *req
r.Model = p.ModelFor(req.Model)
if !p.ModelAvailable(r.Model) {
ok, isProbe := p.ModelSchedulable(r.Model)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
continue
}
resp, err := p.Chat(ctx, &r)
if isProbe {
p.ProbeDone(r.Model)
}
if ctx.Err() != nil {
return nil, "", "", ctx.Err()
}
@ -431,11 +527,15 @@ func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types
p := cands[i]
r := *req
r.Model = p.ModelFor(req.Model)
if !p.ModelAvailable(r.Model) {
ok, isProbe := p.ModelSchedulable(r.Model)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), r.Model)
continue
}
resp, err := p.ChatStream(ctx, &r)
if isProbe {
p.ProbeDone(r.Model)
}
if err == nil {
return resp, p.Name(), r.Model, nil
}
@ -454,11 +554,16 @@ func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.Imag
var lastErr error
for i := 0; i < attempts && i < len(cands); i++ {
p := cands[i]
if !p.ModelAvailable(p.ModelFor(req.Model)) {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), p.ModelFor(req.Model))
im := p.ModelFor(req.Model)
ok, isProbe := p.ModelSchedulable(im)
if !ok {
lastErr = fmt.Errorf("provider %s: model %q cooling down", p.Name(), im)
continue
}
resp, err := p.Image(ctx, req)
if isProbe {
p.ProbeDone(im)
}
if err == nil {
return resp, p.Name(), nil
}

173
internal/scheduler/scheduler_test.go
View File

@ -22,6 +22,14 @@ type fakeProvider struct {
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 {
@ -36,6 +44,25 @@ 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) {
@ -322,3 +349,149 @@ func TestChainResetCooldownAfterSwap(t *testing.T) {
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)
}
}
}