homelab-frontend/internal/proxy/engine_state.go

package proxy

import (
	"bufio"
	"context"
	"fmt"
	"io"
	"math"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

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

// Metric names published by the vLLM engine. Admission is gated on these
// rather than on a per-pod semaphore: every gateway replica reads the same
// engine, so the limit is global by construction (#64).
const (
	metricRequestsRunning = "vllm:num_requests_running"
	metricRequestsWaiting = "vllm:num_requests_waiting"
	metricKVCacheUsage    = "vllm:kv_cache_usage_perc"

	// enginePollInterval is how often each backend's /metrics is scraped.
	enginePollInterval = 2 * time.Second

	// engineStaleThreshold is the number of consecutive failed polls after
	// which the state is marked unhealthy and admission fails open.
	engineStaleThreshold = 3

	// engineMetricsTimeout bounds a single scrape.
	engineMetricsTimeout = 2 * time.Second
)

// EngineState is the gateway's view of one vLLM engine, mirrored from its
// own metrics endpoint.
type EngineState struct {
	Running   int64
	Waiting   int64
	KVUsage   float64
	UpdatedAt time.Time
	Healthy   bool
}

// InFlight returns the number of requests the engine currently holds:
// running plus waiting.
func (s EngineState) InFlight() int64 {
	return s.Running + s.Waiting
}

// EngineObserver polls one vLLM backend's /metrics and publishes the gauges
// admission is gated on. Healthy is false until the first successful poll and
// after engineStaleThreshold consecutive failures.
type EngineObserver struct {
	addr     string
	interval time.Duration
	client   *http.Client

	mu       sync.Mutex
	state    EngineState
	failures int
	degraded bool
	tick     chan struct{}
}

// NewEngineObserver creates an observer for the backend at addr.
func NewEngineObserver(addr string, interval time.Duration) *EngineObserver {
	return &EngineObserver{
		addr:     addr,
		interval: interval,
		// Own client: the proxy transports carry socket options and connect
		// timeouts tuned for inference, not for a short metrics scrape.
		client: &http.Client{Timeout: engineMetricsTimeout},
		tick:   make(chan struct{}),
	}
}

// Start launches the polling goroutine. It returns immediately; the goroutine
// exits when ctx is cancelled.
func (o *EngineObserver) Start(ctx context.Context) {
	go func() {
		ticker := time.NewTicker(o.interval)
		defer ticker.Stop()
		// The first scrape waits one interval rather than firing at Start().
		// Scraping immediately makes the gateway contact every observed backend
		// before any request needs one, which is wasted work and a surprising
		// side effect for anything asserting on upstream traffic. State is
		// unknown until the first poll, and unknown fails open, so the deferral
		// costs at most one interval of ungated admission. The startup
		// self-check (decision Q21) still runs, on that first tick.
		startup := true
		for {
			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
				o.poll(ctx, startup)
				startup = false
			}
		}
	}()
}

// Get returns the most recent engine state.
func (o *EngineObserver) Get() EngineState {
	o.mu.Lock()
	defer o.mu.Unlock()
	return o.state
}

// Tick returns a channel that is closed and replaced after every successful
// poll, waking requests parked on engine admission.
func (o *EngineObserver) Tick() <-chan struct{} {
	o.mu.Lock()
	defer o.mu.Unlock()
	return o.tick
}

// Degraded reports whether the startup self-check found the metrics feed
// unusable (parse failure or missing series).
func (o *EngineObserver) Degraded() bool {
	o.mu.Lock()
	defer o.mu.Unlock()
	return o.degraded
}

// poll scrapes /metrics once. startup marks the first poll, whose failure is a
// self-diagnosis signal rather than a transient blip.
func (o *EngineObserver) poll(ctx context.Context, startup bool) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://"+o.addr+"/metrics", nil)
	if err != nil {
		o.recordFailure()
		logging.Warnf("engine metrics request build failed", map[string]string{
			"address": o.addr,
			"err":     err.Error(),
		})
		return
	}
	resp, err := o.client.Do(req)
	if err != nil {
		o.recordFailure()
		logging.Warnf("engine metrics scrape failed", map[string]string{
			"address": o.addr,
			"err":     err.Error(),
		})
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		o.recordFailure()
		logging.Warnf("engine metrics scrape returned non-200", map[string]string{
			"address": o.addr,
			"status":  strconv.Itoa(resp.StatusCode),
		})
		return
	}

	running, waiting, kvUsage, err := ParseEngineMetrics(resp.Body)
	if err != nil {
		if startup {
			// Startup self-check (decision Q21): metric-name drift must be
			// loud, never a silent fall-back to fail-open admission.
			o.mu.Lock()
			o.degraded = true
			o.mu.Unlock()
			observability.Global().SetObserverDegraded(true)
			logging.Errorf("engine metrics unusable at startup", err, map[string]string{
				"address": o.addr,
			})
		} else {
			logging.Warnf("engine metrics parse failed", map[string]string{
				"address": o.addr,
				"err":     err.Error(),
			})
		}
		o.recordFailure()
		return
	}

	o.recordSuccess(running, waiting, kvUsage)
}

// recordSuccess stores a fresh reading and wakes parked waiters.
func (o *EngineObserver) recordSuccess(running, waiting int64, kvUsage float64) {
	o.mu.Lock()
	o.state = EngineState{
		Running:   running,
		Waiting:   waiting,
		KVUsage:   kvUsage,
		UpdatedAt: time.Now(),
		Healthy:   true,
	}
	o.failures = 0
	close(o.tick)
	o.tick = make(chan struct{})
	o.mu.Unlock()

	m := observability.Global()
	m.SetVLLMRequestsRunning(running)
	m.SetVLLMKVUsagePerc(kvUsage)
}

// recordFailure counts a failed poll and, on the third consecutive failure,
// marks the state unhealthy and counts the staleness transition once — not
// once per failed poll. Last known values are kept.
func (o *EngineObserver) recordFailure() {
	o.mu.Lock()
	o.failures++
	stale := o.failures == engineStaleThreshold
	if o.failures >= engineStaleThreshold {
		o.state.Healthy = false
	}
	failures := o.failures
	o.mu.Unlock()

	if stale {
		observability.Global().IncrEngineStateStale()
		logging.Warnf("engine metrics stale; admission is failing open", map[string]string{
			"address":  o.addr,
			"failures": strconv.Itoa(failures),
		})
	}
}

// ParseEngineMetrics extracts the running, waiting and KV-cache-usage gauges
// from a vLLM Prometheus exposition body. Comment lines, unrelated series and
// non-finite samples are ignored; values are summed across repeated label sets.
// It returns an error naming any of the three series that is absent.
func ParseEngineMetrics(r io.Reader) (running int64, waiting int64, kvUsage float64, err error) {
	scanner := bufio.NewScanner(r)
	// The real exposition body is large; give the scanner room per line.
	scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024)

	var hasRunning, hasWaiting, hasKV bool
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())
		if line == "" || strings.HasPrefix(line, "#") {
			continue
		}
		name, value, ok := splitMetricSample(line)
		if !ok {
			continue
		}
		switch name {
		case metricRequestsRunning, metricRequestsWaiting, metricKVCacheUsage:
		default:
			continue
		}
		f, err := strconv.ParseFloat(value, 64)
		if err != nil || math.IsNaN(f) || math.IsInf(f, 0) {
			continue
		}
		switch name {
		case metricRequestsRunning:
			running += int64(f)
			hasRunning = true
		case metricRequestsWaiting:
			waiting += int64(f)
			hasWaiting = true
		case metricKVCacheUsage:
			kvUsage += f
			hasKV = true
		}
	}
	if err := scanner.Err(); err != nil {
		return 0, 0, 0, fmt.Errorf("read engine metrics: %w", err)
	}

	var missing []string
	if !hasRunning {
		missing = append(missing, metricRequestsRunning)
	}
	if !hasWaiting {
		missing = append(missing, metricRequestsWaiting)
	}
	if !hasKV {
		missing = append(missing, metricKVCacheUsage)
	}
	if len(missing) > 0 {
		return 0, 0, 0, fmt.Errorf("engine metrics missing series: %s", strings.Join(missing, ", "))
	}
	return running, waiting, kvUsage, nil
}

// splitMetricSample splits one Prometheus exposition line into its metric name
// and its first sample value. Lines without a value are reported as not ok.
func splitMetricSample(line string) (name, value string, ok bool) {
	end := strings.IndexAny(line, "{ \t")
	if end <= 0 {
		return "", "", false
	}
	name = line[:end]
	rest := line[end:]
	if rest[0] == '{' {
		closeBrace := strings.LastIndex(rest, "}")
		if closeBrace < 0 {
			return "", "", false
		}
		rest = rest[closeBrace+1:]
	}
	fields := strings.Fields(rest)
	if len(fields) == 0 {
		return "", "", false
	}
	return name, fields[0], true
}