homelab-frontend/internal/proxy/scheduler.go

package proxy

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

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

// 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)
type Backend struct {
	Address string
	queue   chan struct{}
	slots   chan struct{}
}

// 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.
func (b *Backend) Acquire(ctx context.Context) error {
	// 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.
func NewScheduler(addresses []string, maxSlots, maxQueue int) *Scheduler {
	backends := make([]*Backend, len(addresses))
	for i, addr := range addresses {
		backends[i] = newBackend(addr, maxSlots, maxQueue)
	}
	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)
}