llmgw/internal/router/admission/controller.go

package admission

import (
	"context"
	"errors"
	"fmt"
	"slices"
	"sync"
	"time"
)

// Band is a priority class. Requests in a band are held once pool saturation
// exceeds the band's ceiling, so the cluster degrades from the bottom up:
// batch traffic slows first and interactive traffic keeps its latency.
type Band struct {
	Name string `koanf:"name"`
	// Priority orders bands; higher is served first.
	Priority int `koanf:"priority"`
	// SaturationCeiling is the pool saturation above which this band waits.
	// 1.0 means "never hold this band back".
	SaturationCeiling float64 `koanf:"saturation_ceiling"`
}

// Config configures the controller.
type Config struct {
	Bands       []Band `koanf:"bands"`
	DefaultBand string `koanf:"default_band"`

	Thresholds Thresholds `koanf:"-"`

	// QueueTTL drops a request that has waited this long. Sending the GPU work
	// the client has already given up on helps nobody.
	QueueTTL time.Duration `koanf:"queue_ttl"`
	// QueueCapacityPerBand bounds each band's queue. Rejecting immediately
	// beats accumulating silently.
	QueueCapacityPerBand int `koanf:"queue_capacity_per_band"`
	// QueueCapacityTotal bounds everything queued, across bands. When it is
	// reached, a request is not refused while a lower band has requests
	// queued: the newest of the lowest band is turned away instead
	// (displacement), since the oldest have waited longest and are nearest
	// their TTL. Zero takes half of every band's capacity together.
	QueueCapacityTotal int `koanf:"queue_capacity_total"`
	// Interval is how often the dispatcher re-checks saturation. Releases also
	// wake it, so this is only the backstop for saturation falling on its own.
	Interval time.Duration `koanf:"interval"`
}

// DefaultBands returns the four-band scheme from plan §7.2.2.
func DefaultBands() []Band {
	return []Band{
		{Name: "platinum", Priority: 400, SaturationCeiling: 1.00},
		{Name: "gold", Priority: 300, SaturationCeiling: 0.85},
		{Name: "silver", Priority: 200, SaturationCeiling: 0.60},
		{Name: "bronze", Priority: 100, SaturationCeiling: 0.30},
	}
}

// DefaultConfig returns a usable configuration.
func DefaultConfig() Config {
	return Config{
		Bands:                DefaultBands(),
		DefaultBand:          "gold",
		Thresholds:           Thresholds{TargetWaitingPerReplica: 4, KVCeiling: 0.9},
		QueueTTL:             30 * time.Second,
		QueueCapacityPerBand: 2000,
		Interval:             10 * time.Millisecond,
	}
}

func (c *Config) normalize() error {
	if len(c.Bands) == 0 {
		c.Bands = DefaultBands()
	}
	slices.SortFunc(c.Bands, func(a, b Band) int { return b.Priority - a.Priority })

	seen := map[string]struct{}{}
	for i := range c.Bands {
		b := &c.Bands[i]
		if b.Name == "" {
			return fmt.Errorf("bands[%d]: name is required", i)
		}
		if _, dup := seen[b.Name]; dup {
			return fmt.Errorf("band %q is declared twice", b.Name)
		}
		seen[b.Name] = struct{}{}
		if b.SaturationCeiling <= 0 {
			b.SaturationCeiling = 1
		}
	}
	if c.DefaultBand == "" {
		c.DefaultBand = c.Bands[len(c.Bands)-1].Name
	}
	if _, ok := seen[c.DefaultBand]; !ok {
		return fmt.Errorf("default_band %q is not one of the configured bands", c.DefaultBand)
	}
	c.Thresholds.normalize()
	if c.QueueTTL <= 0 {
		c.QueueTTL = 30 * time.Second
	}
	if c.QueueCapacityPerBand <= 0 {
		c.QueueCapacityPerBand = 2000
	}
	if c.QueueCapacityTotal <= 0 {
		c.QueueCapacityTotal = max(c.QueueCapacityPerBand*len(c.Bands)/2, 1)
	}
	if c.Interval <= 0 {
		c.Interval = 10 * time.Millisecond
	}
	return nil
}

// SaturationFunc reports a model's pool saturation and whether it is known.
//
// outstanding is how many requests this controller has admitted and not yet
// seen released, so the caller can fold in load the engines have not reported.
type SaturationFunc func(model string, outstanding int) (float64, bool)

// Request identifies who is asking for what.
type Request struct {
	Model string
	// Band selects the priority class; unknown names fall back to the default.
	Band string
	// Tenant is the fairness key within a band. Two requests sharing a tenant
	// compete with each other rather than with everyone else.
	Tenant string
}

var (
	// ErrQueueFull means the band's queue is at capacity.
	ErrQueueFull = errors.New("admission queue is full")
	// ErrQueueTimeout means the request waited longer than the TTL.
	ErrQueueTimeout = errors.New("admission queue timeout")
	// ErrClosed means the controller is shutting down.
	ErrClosed = errors.New("admission controller closed")
	// ErrDisplaced means a queued request was turned away to make room for
	// one of a higher band. Worth retrying: the queue is full now, not for
	// good.
	ErrDisplaced = errors.New("displaced from the admission queue by higher-priority requests")
)

// Stats is a snapshot of one band's queue, for metrics.
type Stats struct {
	Band  string
	Depth int
}

// Controller admits requests according to band priority and tenant fairness.
type Controller struct {
	cfg    Config
	satFn  SaturationFunc
	bandOf map[string]int

	mu     sync.Mutex
	queues []*bandQueue
	closed bool

	wake chan struct{}
	done chan struct{}
	stop chan struct{}
	once sync.Once

	// outstanding counts admitted-but-unreleased requests per model. It is the
	// only load signal available with zero lag, which is what stops a
	// simultaneous burst from passing the check en masse.
	outstanding map[string]int

	// Observed counters, read through Metrics.
	admitted  map[string]int64
	rejected  map[string]int64
	waitTotal map[string]time.Duration
}

// New builds a controller. Start must be called before it will admit anything
// that has to wait.
func New(cfg Config, satFn SaturationFunc) (*Controller, error) {
	if err := cfg.normalize(); err != nil {
		return nil, err
	}
	c := &Controller{
		cfg:         cfg,
		satFn:       satFn,
		bandOf:      make(map[string]int, len(cfg.Bands)),
		wake:        make(chan struct{}, 1),
		done:        make(chan struct{}),
		stop:        make(chan struct{}),
		outstanding: map[string]int{},
		admitted:    map[string]int64{},
		rejected:    map[string]int64{},
		waitTotal:   map[string]time.Duration{},
	}
	for i, b := range cfg.Bands {
		c.bandOf[b.Name] = i
		c.queues = append(c.queues, newBandQueue(b))
	}
	return c, nil
}

// Bands returns the configured bands, highest priority first.
func (c *Controller) Bands() []Band { return slices.Clone(c.cfg.Bands) }

// Start launches the dispatcher.
func (c *Controller) Start() {
	go c.dispatchLoop()
}

// Close stops the dispatcher and fails everything still queued.
func (c *Controller) Close() {
	c.once.Do(func() {
		close(c.stop)
		<-c.done
	})
}

// Lease is held for the lifetime of an admitted request.
type Lease struct {
	c     *Controller
	band  string
	model string
	once  sync.Once
}

// Release returns the slot and wakes the dispatcher, since finishing a request
// is the most likely reason saturation just dropped.
func (l *Lease) Release() {
	if l == nil || l.c == nil {
		return
	}
	l.once.Do(func() {
		l.c.mu.Lock()
		if n := l.c.outstanding[l.model]; n > 1 {
			l.c.outstanding[l.model] = n - 1
		} else {
			delete(l.c.outstanding, l.model)
		}
		l.c.mu.Unlock()
		l.c.signal()
	})
}

// Band reports which band admitted the request.
func (l *Lease) Band() string {
	if l == nil {
		return ""
	}
	return l.band
}

// Acquire blocks until the request may proceed, the context ends, or the TTL
// expires.
func (c *Controller) Acquire(ctx context.Context, r Request) (*Lease, error) {
	idx, ok := c.bandOf[r.Band]
	if !ok {
		idx = c.bandOf[c.cfg.DefaultBand]
	}
	band := c.cfg.Bands[idx]

	// Fast path: under the ceiling with nobody already waiting in this band.
	// Skipping the queue here is what keeps the common case free of latency.
	c.mu.Lock()
	if c.closed {
		c.mu.Unlock()
		return nil, ErrClosed
	}
	if c.queues[idx].len() == 0 && c.admits(band, r.Model) {
		c.grant(band.Name, r.Model)
		c.mu.Unlock()
		return &Lease{c: c, band: band.Name, model: r.Model}, nil
	}

	if c.queues[idx].len() >= c.cfg.QueueCapacityPerBand {
		c.rejected[band.Name]++
		c.mu.Unlock()
		return nil, fmt.Errorf("%w: band %q", ErrQueueFull, band.Name)
	}
	if c.queuedLocked() >= c.cfg.QueueCapacityTotal && !c.displaceBelowLocked(idx) {
		c.rejected[band.Name]++
		c.mu.Unlock()
		return nil, fmt.Errorf("%w: %d queued in all, and nothing of a lower band to make room",
			ErrQueueFull, c.cfg.QueueCapacityTotal)
	}

	w := &waiter{
		tenant:   tenantKey(r),
		model:    r.Model,
		ready:    make(chan struct{}),
		enqueued: time.Now(),
		deadline: time.Now().Add(c.cfg.QueueTTL),
	}
	c.queues[idx].push(w)
	c.mu.Unlock()

	c.signal()

	select {
	case <-w.ready:
		c.mu.Lock()
		if c.closed {
			c.mu.Unlock()
			return nil, ErrClosed
		}
		if w.err != nil {
			c.mu.Unlock()
			return nil, w.err
		}
		c.waitTotal[band.Name] += time.Since(w.enqueued)
		c.mu.Unlock()
		return &Lease{c: c, band: band.Name, model: r.Model}, nil

	case <-ctx.Done():
		c.remove(idx, w, "client_gone")
		return nil, ctx.Err()

	case <-time.After(time.Until(w.deadline)):
		c.remove(idx, w, "ttl")
		return nil, fmt.Errorf("%w after %v", ErrQueueTimeout, c.cfg.QueueTTL)
	}
}

// remove takes a waiter out of its queue unless it was admitted first.
func (c *Controller) remove(idx int, w *waiter, reason string) {
	c.mu.Lock()
	defer c.mu.Unlock()
	if w.admitted || w.err != nil {
		// It won the race, or was displaced; the caller will see ready
		// closed next time round.
		return
	}
	c.queues[idx].remove(w)
	c.rejected[c.cfg.Bands[idx].Name+"/"+reason]++
}

// grant records an admission. The caller holds the lock.
func (c *Controller) grant(band, model string) {
	c.admitted[band]++
	c.outstanding[model]++
}

// admits reports whether a band may proceed right now. The caller holds the lock.
func (c *Controller) admits(b Band, model string) bool {
	sat, known := c.satFn(model, c.outstanding[model])
	if !known {
		// Nothing is known about this model's pool; let it through so dispatch
		// can return a precise error instead of a queue timeout.
		return true
	}
	return sat <= b.SaturationCeiling
}

func (c *Controller) signal() {
	select {
	case c.wake <- struct{}{}:
	default:
	}
}

func (c *Controller) dispatchLoop() {
	defer close(c.done)
	t := time.NewTicker(c.cfg.Interval)
	defer t.Stop()

	for {
		select {
		case <-c.stop:
			c.drain()
			return
		case <-c.wake:
			c.dispatch()
		case <-t.C:
			c.dispatch()
		}
	}
}

// dispatch walks the bands from highest priority down, admitting whatever the
// current saturation allows.
func (c *Controller) dispatch() {
	now := time.Now()

	c.mu.Lock()
	defer c.mu.Unlock()

	for i, q := range c.queues {
		q.dropExpired(now, func() { c.rejected[c.cfg.Bands[i].Name+"/ttl"]++ })

		band := c.cfg.Bands[i]
		// Each pass admits at most the queue length, so a band cannot spin
		// forever while saturation is recomputed from stale metrics.
		for range q.len() {
			w := q.peek()
			if w == nil || !c.admits(band, w.model) {
				break
			}
			q.pop()
			w.admitted = true
			c.grant(band.Name, w.model)
			close(w.ready)
		}
	}
}

func (c *Controller) drain() {
	c.mu.Lock()
	defer c.mu.Unlock()
	c.closed = true
	for _, q := range c.queues {
		for {
			w := q.pop()
			if w == nil {
				break
			}
			// Closing without marking admitted lets Acquire report ErrClosed.
			close(w.ready)
			w.admitted = true
		}
	}
}

// Depths returns each band's current queue depth.
func (c *Controller) Depths() []Stats {
	c.mu.Lock()
	defer c.mu.Unlock()
	out := make([]Stats, 0, len(c.queues))
	for i, q := range c.queues {
		out = append(out, Stats{Band: c.cfg.Bands[i].Name, Depth: q.len()})
	}
	return out
}

// tenantKey picks the fairness key, falling back to the band so that unknown
// tenants at least compete as one group rather than each looking unique.
func tenantKey(r Request) string {
	if r.Tenant != "" {
		return r.Tenant
	}
	return "anonymous"
}

// queuedLocked is how many requests wait in every band. Called with mu held.
func (c *Controller) queuedLocked() int {
	n := 0
	for _, q := range c.queues {
		n += q.len()
	}
	return n
}

// displaceBelowLocked turns away the newest request of the lowest band below
// idx that has any queued, and says whether there was one. Called with mu held.
func (c *Controller) displaceBelowLocked(idx int) bool {
	// Bands are sorted highest priority first, so lower bands have larger
	// indices, and the lowest is last.
	for j := len(c.queues) - 1; j > idx; j-- {
		w := c.queues[j].newest()
		if w == nil {
			continue
		}
		c.queues[j].remove(w)
		w.err = fmt.Errorf("%w (band %q)", ErrDisplaced, c.cfg.Bands[j].Name)
		close(w.ready)
		c.rejected[c.cfg.Bands[j].Name+"/displaced"]++
		return true
	}
	return false
}