ModelRouter/internal/scheduler/scheduler.go

// Package scheduler implements request scheduling across providers: direct
// fallback scheduling over candidate lists, and the AUTO chain (tiers with
// per-tier round-robin cursors, preference ordering, token-quota windows and
// per-(source,model) cooldown awareness) per the target architecture in
// plan.md.
package scheduler

import (
	"context"
	"errors"
	"fmt"
	"sort"
	"strings"
	"sync/atomic"
	"time"

	"llmsproxy/internal/types"
)

// busyWait is how long a fully-busy tier is polled for a free slot before the
// request falls through to the next tier (bounded wait, plan 2.3).
var busyWait = 2 * time.Second

// busyPoll is the polling interval while waiting for a busy tier.
var busyPoll = 100 * time.Millisecond

// Scheduler drives one chat tool call across the candidate provider chain.
type Scheduler struct {
	// MaxRetries how many fallback providers to try before failing.
	MaxRetries int
}

func New(maxRetries int) *Scheduler {
	if maxRetries < 0 {
		maxRetries = 0
	}
	return &Scheduler{MaxRetries: maxRetries}
}

// Provider is the minimal interface the scheduler needs to schedule over.
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)
	Image(ctx context.Context, req *types.ImageGenRequest) (*types.UnifiedResponse, error)
}

// ---- AUTO chain ----

// Rule is one persisted AUTO chain slot (mirror of config.ModelScope).
type Rule struct {
	Model  string
	Source string
	Tier   int
	Quota  int64
	Period string
	Hours  int64
}

// Slot is one schedulable chain position: a model pinned to its provider,
// with an optional token-quota window. Slots are immutable after build.
type Slot struct {
	Model  string
	Source string
	Quota  int64
	Period string
	Hours  int64
	Prov   Provider
}

// TierNode is one priority tier. Slots keep their configured order (the
// stable base for preference ordering). next is the round-robin cursor: it
// holds the last used slot index (-1 = none yet), so the very first request
// starts at the configured order and later ones rotate.
type TierNode struct {
	Tier  int
	Slots []*Slot
	next  atomic.Int64
}

// NextStart advances the tier cursor and returns the start index for the next
// scheduling run (first run: index 0).
func (tn *TierNode) NextStart() int64 {
	return tn.next.Add(1)
}

// Chain is the immutable AUTO scheduling plan. A rebuilt chain is swapped in
// atomically; per-tier cursors live inside the chain and are shared across
// requests (rotation state resets when the chain is rebuilt, e.g. after
// editing the rules — acceptable, the swap also resets cooldowns).
type Chain struct {
	Tiers []*TierNode // ascending tier order (tier 1 = highest priority, tried first)
}

// TierErrors is the per-tier failure summary carried by ChainErr. Errors
// (TierError or skipped-tier reasons) are collected in tier order.
type TierError struct {
	Tier   int
	Source string
	Model  string
	Err    error
}

// ChainErr is returned by chain scheduling when every AUTO tier failed. Its
// message summarizes each failed tier (which source/model and why) so a 503
// names the culprits instead of the bare "no provider available".
type ChainErr struct {
	Tiers   []TierError
	Skipped []string // whole-tier reasons (cooling / quota / all busy)
}

func (e *ChainErr) Error() string {
	var b strings.Builder
	b.WriteString("all auto tiers failed: ")
	first := true
	for _, t := range e.Tiers {
		if !first {
			b.WriteString("; ")
		}
		first = false
		fmt.Fprintf(&b, "tier %d %s/%s: %v", t.Tier, t.Source, t.Model, t.Err)
	}
	for _, s := range e.Skipped {
		if !first {
			b.WriteString("; ")
		}
		first = false
		b.WriteString(s)
	}
	return b.String()
}

// BuildChain groups rules into ascending tiers (tier 1 = highest priority,
// tried first) and resolves each slot's provider via prov. Rules whose
// provider resolves to nil are dropped (the source no longer serves the
// model). Slot order within a tier follows the configured rule order.
func BuildChain(rules []Rule, prov func(model, source string) Provider) *Chain {
	byTier := map[int][]*Slot{}
	var tiers []int
	for _, r := range rules {
		p := prov(r.Model, r.Source)
		if p == nil {
			continue
		}
		if _, ok := byTier[r.Tier]; !ok {
			tiers = append(tiers, r.Tier)
		}
		byTier[r.Tier] = append(byTier[r.Tier], &Slot{
			Model:  r.Model,
			Source: r.Source,
			Quota:  r.Quota,
			Period: r.Period,
			Hours:  r.Hours,
			Prov:   p,
		})
	}
	sort.Slice(tiers, func(i, j int) bool { return tiers[i] < tiers[j] })
	ch := &Chain{}
	for _, t := range tiers {
		tn := &TierNode{Tier: t, Slots: byTier[t]}
		tn.next.Store(-1)
		ch.Tiers = append(ch.Tiers, tn)
	}
	return ch
}

// tierResult is the outcome of one scheduling run over one tier.
type tierResult struct {
	resp   *types.UnifiedResponse
	chunks <-chan types.UnifiedChunk
	src    string
	model  string
	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.
//
// 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++ {
		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
		r.Model = sl.Model
		if stream {
			chunks, err := sl.Prov.ChatStream(ctx, &r)
			if err == nil {
				return tierResult{chunks: chunks, src: sl.Source, model: sl.Model}
			}
			if ctx.Err() != nil {
				return tierResult{}
			}
			if errors.Is(err, types.ErrBusy) {
				continue
			}
			hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
			continue
		}
		resp, err := sl.Prov.Chat(ctx, &r)
		if err == nil {
			return tierResult{resp: resp, src: sl.Source, model: sl.Model}
		}
		if ctx.Err() != nil {
			return tierResult{}
		}
		if errors.Is(err, types.ErrBusy) {
			continue
		}
		hard = append(hard, TierError{Tier: tn.Tier, Source: sl.Source, Model: sl.Model, Err: err})
	}
	return tierResult{hard: hard}
}

// chainDrive runs a request down the chain (plan 2.3): tiers ascending (tier
// 1, the highest priority, first), per-tier round-robin starting at the tier
// cursor, same-tier runs ordered by preference (negative prefs sink but stay
// reachable). Quota-exhausted and cooling slots are filtered up front; a
// fully busy tier is polled for a bounded time before falling through.
// Failures are summarized in *ChainErr for the caller to map to HTTP 503.
func (s *Scheduler) chainDrive(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool, stream bool) (*types.UnifiedResponse, <-chan types.UnifiedChunk, string, string, error) {
	if chain == nil || len(chain.Tiers) == 0 {
		return nil, nil, "", "", fmt.Errorf("no auto slot configured")
	}
	var ce ChainErr
	for _, tn := range chain.Tiers {
		// 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 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
		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
			}
			return nil, nil, "", "", err
		}
		releaseProbes(cands)
	}
	if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
		return nil, nil, "", "", fmt.Errorf("no auto slot configured")
	}
	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
// summarizing every tier.
func (s *Scheduler) ChainChat(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool) (*types.UnifiedResponse, string, string, error) {
	resp, _, src, model, err := s.chainDrive(ctx, chain, req, exhausted, false)
	return resp, src, model, err
}

// ChainChatStream runs a streaming AUTO request down the chain. A slot is
// abandoned only on connect failures / busy (before its first chunk); after a
// stream starts it is pinned. Same return contract as ChainChat.
func (s *Scheduler) ChainChatStream(ctx context.Context, chain *Chain, req *types.ChatRequest, exhausted func(*Slot) bool) (<-chan types.UnifiedChunk, string, string, error) {
	_, chunks, src, model, err := s.chainDrive(ctx, chain, req, exhausted, true)
	return chunks, src, model, err
}

// ChainImage runs an image-generation AUTO request down the chain: tiers
// ascending (tier 1 highest priority), per-tier round-robin, same-tier order
// by preference. Each slot's model is pinned to its own image id (ModelFor),
// so a fallback switches per source. Cooling-down slots are skipped. Returns
// the response, serving source and the exact model id used; on total failure
// a *ChainErr summarizing every tier.
func (s *Scheduler) ChainImage(ctx context.Context, chain *Chain, req *types.ImageGenRequest) (*types.UnifiedResponse, string, string, error) {
	if chain == nil || len(chain.Tiers) == 0 {
		return nil, "", "", fmt.Errorf("no image auto slot configured")
	}
	var ce ChainErr
	for _, tn := range chain.Tiers {
		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++ {
			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) {
				continue
			}
			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...)
		}
	}
	if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
		return nil, "", "", fmt.Errorf("no image auto slot configured")
	}
	return nil, "", "", &ce
}

// ---- direct scheduling ----

// Chat runs a chat request across cands, falling back on failure. Each
// candidate receives a request pinned to its own model (ModelFor), so a
// fallback switches the model id per provider instead of reusing the first
// candidate's model name. On success it returns the response together with
// the name of the provider and the exact model id that served the request.
// A candidate whose (source, model) pair is cooling down is skipped like a
// busy one — direct paths share the "cooldown is the only hard skip"
// semantics of the AUTO chain (plan 2.3/2.5); otherwise persistent direct
// traffic would keep renewing a capped auth cooldown forever.
func (s *Scheduler) Chat(ctx context.Context, cands []Provider, req *types.ChatRequest) (*types.UnifiedResponse, string, string, error) {
	attempts := s.MaxRetries + 1
	var lastErr error
	for i := 0; i < attempts && i < len(cands); i++ {
		p := cands[i]
		r := *req
		r.Model = p.ModelFor(req.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()
		}
		if err == nil {
			return resp, p.Name(), r.Model, nil
		}
		lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
	}
	if lastErr == nil {
		if len(cands) == 0 {
			lastErr = fmt.Errorf("no provider available")
		}
	}
	return nil, "", "", lastErr
}

// ChatStream runs a streaming chat across cands, falling back early on
// connect errors. The request model is pinned per candidate like Chat. On
// success it returns the chunk channel plus the serving provider name and
// model id.
func (s *Scheduler) ChatStream(ctx context.Context, cands []Provider, req *types.ChatRequest) (<-chan types.UnifiedChunk, string, string, error) {
	attempts := s.MaxRetries + 1
	var lastErr error
	for i := 0; i < attempts && i < len(cands); i++ {
		p := cands[i]
		r := *req
		r.Model = p.ModelFor(req.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
		}
		lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
	}
	if lastErr == nil && len(cands) == 0 {
		lastErr = fmt.Errorf("no provider available")
	}
	return nil, "", "", lastErr
}

// Image runs an image-generation request across cands; returns the used
// provider name on success.
func (s *Scheduler) Image(ctx context.Context, cands []Provider, req *types.ImageGenRequest) (*types.UnifiedResponse, string, error) {
	attempts := s.MaxRetries + 1
	var lastErr error
	for i := 0; i < attempts && i < len(cands); i++ {
		p := cands[i]
		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
		}
		lastErr = fmt.Errorf("provider %s: %w", p.Name(), err)
	}
	if lastErr == nil && len(cands) == 0 {
		lastErr = fmt.Errorf("no provider available")
	}
	return nil, "", lastErr
}