homelab-frontend/internal/proxy/scheduler.go

package proxy

import (
	"context"
	"errors"
	"hash/fnv"
	"sort"
	"sync"
	"sync/atomic"
	"time"

	"forgejo.riotpiao.com/rock/homelab-frontend/internal/observability"
)

// ErrQueueFull is returned when the in-memory queue is at capacity.
var ErrQueueFull = errors.New("backend queue full")

// ErrAdmissionTimeout is returned when a request parked on engine admission
// exceeds the configured maximum park duration.
var ErrAdmissionTimeout = errors.New("admission timeout")

// AdmissionLimits are the engine-wide admission thresholds applied to every
// backend of a scheduler (#64).
type AdmissionLimits struct {
	MaxInFlight int
	MaxKVUsage  float64
	MaxPark     time.Duration
}

// Backend represents a single vLLM upstream with a bounded in-memory queue.
// Two channels model the queue:
//
//	queue  — waiting requests (bounded by VLLM_MAX_QUEUE_DEPTH)
//	slots  — concurrent in-service requests (bounded by VLLM_MAX_NUM_SEQS)
//
// The engine gate is applied before the queue: admission reads the engine's
// own published gauges via observer, which is global across replicas.
type Backend struct {
	Address     string
	queue       chan struct{}
	slots       chan struct{}
	observer    *EngineObserver
	maxInFlight int
	maxKVUsage  float64
	maxPark     time.Duration
}

// newBackend creates a Backend with given queue depth and concurrency limit.
func newBackend(addr string, maxSlots, maxQueue int) *Backend {
	slots := make(chan struct{}, maxSlots)
	for i := 0; i < maxSlots; i++ {
		slots <- struct{}{}
	}
	return &Backend{
		Address: addr,
		queue:   make(chan struct{}, maxQueue),
		slots:   slots,
	}
}

// Acquire enters the in-memory queue then waits for a vLLM concurrency slot.
// Returns ErrQueueFull immediately if the queue channel is at capacity (no goroutine queued).
// Returns ctx.Err() if the context is cancelled while waiting for a slot.
// Returns ErrAdmissionTimeout if the engine gate parks the request past maxPark.
func (b *Backend) Acquire(ctx context.Context) error {
	return b.acquire(ctx, nil)
}

// acquire is Acquire with an optional onPark callback. onPark is invoked once,
// on the calling goroutine, the first time the request has to wait for engine
// admission; the parked streaming path uses it to start an SSE heartbeat.
func (b *Backend) acquire(ctx context.Context, onPark func()) error {
	// Engine gate (#64): admit on the engine's own gauges, not on this pod's
	// counters. Failing OPEN on !Healthy is deliberate — a metrics outage must
	// not become an inference outage; the three-replica gate is only a fairness
	// improvement on top of the engine's own queueing.
	if b.observer != nil {
		deadline := time.Now().Add(b.maxPark)
		timer := time.NewTimer(time.Until(deadline))
		defer timer.Stop()
		parked := false
		for {
			// State and tick channel are read together under one lock: taking
			// them separately can miss a poll that lands between the two calls.
			st, tick := b.observer.Snapshot()
			if !st.Healthy || (st.InFlight() < int64(b.maxInFlight) && st.KVUsage < b.maxKVUsage) {
				break
			}
			if !parked {
				observability.Global().IncrAdmissionParked()
				parked = true
				if onPark != nil {
					onPark()
				}
			}
			select {
			case <-tick:
			case <-ctx.Done():
				return ctx.Err()
			case <-timer.C:
				return ErrAdmissionTimeout
			}
		}
	}

	// Non-blocking queue entry: reject immediately when queue is full.
	select {
	case b.queue <- struct{}{}:
		// entered queue
	default:
		return ErrQueueFull
	}
	// Wait for a concurrency slot, then leave the queue.
	select {
	case <-b.slots:
		<-b.queue // leave queue; now in service
		return nil
	case <-ctx.Done():
		<-b.queue // leave queue; cancelled
		return ctx.Err()
	}
}

// Release returns a concurrency slot to the pool after the upstream request completes.
func (b *Backend) Release() {
	b.slots <- struct{}{}
}

// Scheduler routes LLM requests to backends using rendezvous hashing (cache-aware)
// or round-robin (no cache key). Failed backends are excluded for a configurable window.
type Scheduler struct {
	mu          sync.RWMutex
	backends    []*Backend
	failedUntil map[string]time.Time
	failWindow  time.Duration
	rrIdx       atomic.Uint64
}

// NewScheduler creates a Scheduler for the given addresses.
// maxSlots = VLLM_MAX_NUM_SEQS, maxQueue = VLLM_MAX_QUEUE_DEPTH.
// observers holds one engine observer per backend address (may be nil).
func NewScheduler(addresses []string, maxSlots, maxQueue int, observers map[string]*EngineObserver, limits AdmissionLimits) *Scheduler {
	backends := make([]*Backend, len(addresses))
	for i, addr := range addresses {
		backends[i] = newBackend(addr, maxSlots, maxQueue)
		backends[i].observer = observers[addr]
		backends[i].maxInFlight = limits.MaxInFlight
		backends[i].maxKVUsage = limits.MaxKVUsage
		backends[i].maxPark = limits.MaxPark
	}
	return &Scheduler{
		backends:    backends,
		failedUntil: make(map[string]time.Time),
		failWindow:  30 * time.Second,
	}
}

// Select returns a backend. Uses rendezvous hashing when cacheKey is non-empty
// (same key always routes to the same backend, skipping unhealthy ones).
// Uses round-robin when cacheKey is empty.
func (s *Scheduler) Select(cacheKey string) *Backend {
	s.mu.RLock()
	defer s.mu.RUnlock()

	if len(s.backends) == 0 {
		return nil
	}
	if len(s.backends) == 1 {
		return s.backends[0]
	}
	if cacheKey != "" {
		return s.selectRendezvous(cacheKey)
	}
	return s.selectRoundRobin()
}

// selectRendezvous picks the healthy backend with the highest FNV32a score
// for the key (highest random weight / rendezvous hashing).
// Falls back to next-highest score when primary is excluded by MarkFailed.
// Must be called with mu held.
func (s *Scheduler) selectRendezvous(cacheKey string) *Backend {
	type scored struct {
		b     *Backend
		score uint32
	}
	now := time.Now()
	scores := make([]scored, len(s.backends))
	h := fnv.New32a()
	for i, b := range s.backends {
		h.Reset()
		h.Write([]byte(cacheKey + ":" + b.Address))
		scores[i] = scored{b: b, score: h.Sum32()}
	}
	sort.Slice(scores, func(i, j int) bool {
		return scores[i].score > scores[j].score
	})
	for _, sc := range scores {
		_, failed := s.failedUntil[sc.b.Address]
		if !failed {
			return sc.b
		}
		expiry := s.failedUntil[sc.b.Address]
		if now.After(expiry) {
			return sc.b
		}
	}
	// All backends unhealthy: return highest-scored (fail-open).
	return scores[0].b
}

// selectRoundRobin distributes requests evenly when no cache key is present.
// Must be called with mu held.
func (s *Scheduler) selectRoundRobin() *Backend {
	idx := s.rrIdx.Add(1) - 1
	return s.backends[idx%uint64(len(s.backends))]
}

// MarkFailed excludes addr from routing for the fail window (default 30 s).
// Called by the proxy ErrorHandler on upstream connection failure.
func (s *Scheduler) MarkFailed(addr string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.failedUntil[addr] = time.Now().Add(s.failWindow)
}