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
	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)
}

// 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 {
	n := len(cands)
	var hard []TierError
	for i := 0; i < n; i++ {
		sl := cands[(int(base)+i)%n]
		if !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 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)
		}
		if len(cands) == 0 {
			ce.Skipped = append(ce.Skipped, fmt.Sprintf("tier %d: no schedulable slot (cooling or quota exhausted)", tn.Tier))
			continue
		}
		// preference orders a same-tier run; stable so equal prefs keep order
		sort.SliceStable(cands, func(i, j int) bool {
			return cands[i].Prov.Pref(cands[i].Model) > cands[j].Prov.Pref(cands[j].Model)
		})
		base := tn.NextStart()
		res := runTier(ctx, tn, cands, 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...)
			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)
		for {
			select {
			case <-ctx.Done():
				return nil, nil, "", "", ctx.Err()
			case <-time.After(busyPoll):
			}
			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
			}
		}
	}
	if len(ce.Tiers) == 0 && len(ce.Skipped) == 0 {
		return nil, nil, "", "", fmt.Errorf("no auto slot configured")
	}
	return nil, nil, "", "", &ce
}

// 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
}

// ---- 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.
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)
		resp, err := p.Chat(ctx, &r)
		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)
		resp, err := p.ChatStream(ctx, &r)
		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]
		resp, err := p.Image(ctx, req)
		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
}