llmgw/internal/router/admin.go

package router

import (
	"net/http"
	"net/http/pprof"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/collectors"
	"github.com/prometheus/client_golang/prometheus/promhttp"
)

// AdminHandler returns the diagnostics surface: metrics, health and pprof.
//
// It is served on a separate listener so the client-facing port can be exposed
// without also exposing pprof (plan §9.1). The control plane's management API
// is mounted here too, by the caller.
// RegisterMetrics adds a collector to the /metrics the admin listener serves.
//
// Exists so the control plane can publish what the cluster is made of on the
// same endpoint as the traffic figures. Correlating "the cluster was full" with
// "requests were queueing" is the first thing anyone does after an incident,
// and needing two scrape targets to do it is a needless obstacle.
//
// Called before AdminHandler, because the registry is built there.
func (r *Router) RegisterMetrics(c prometheus.Collector) error {
	r.extraMetrics = append(r.extraMetrics, c)
	return nil
}

func (r *Router) AdminHandler() (*http.ServeMux, error) {
	reg := prometheus.NewRegistry()
	reg.MustRegister(
		collectors.NewGoCollector(),
		collectors.NewProcessCollector(collectors.ProcessCollectorOpts{}),
	)
	for _, c := range r.m.collectors() {
		if err := reg.Register(c); err != nil {
			return nil, err
		}
	}
	if err := reg.Register(newEndpointCollector(r)); err != nil {
		return nil, err
	}
	if err := reg.Register(newKVCollector(r.kv)); err != nil {
		return nil, err
	}
	for _, c := range r.extraMetrics {
		if err := reg.Register(c); err != nil {
			return nil, err
		}
	}

	// Kept, so anything already scraping the admin listener keeps working.
	r.metricsHandler = promhttp.HandlerFor(reg, promhttp.HandlerOpts{})

	mux := http.NewServeMux()
	mux.Handle("GET /metrics", r.metricsHandler)
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	mux.HandleFunc("GET /debug/pprof/", pprof.Index)
	mux.HandleFunc("GET /debug/pprof/cmdline", pprof.Cmdline)
	mux.HandleFunc("GET /debug/pprof/profile", pprof.Profile)
	mux.HandleFunc("GET /debug/pprof/symbol", pprof.Symbol)
	mux.HandleFunc("GET /debug/pprof/trace", pprof.Trace)
	return mux, nil
}

// endpointCollector reports live endpoint state at scrape time rather than
// mirroring it into gauges on every request; the registry already holds it.
type endpointCollector struct {
	r *Router

	inflight *prometheus.Desc
	waiting  *prometheus.Desc
	running  *prometheus.Desc
	kv       *prometheus.Desc
	stale    *prometheus.Desc
	drain    *prometheus.Desc
	breaker  *prometheus.Desc

	queueDepth *prometheus.Desc
	saturation *prometheus.Desc
}

func newEndpointCollector(r *Router) *endpointCollector {
	lbl := []string{"replica", "pool"}
	return &endpointCollector{
		r:        r,
		inflight: prometheus.NewDesc("llmgw_endpoint_inflight", "Requests the router has dispatched and not yet finished.", lbl, nil),
		waiting:  prometheus.NewDesc("llmgw_endpoint_waiting", "Engine queue depth, as last scraped.", lbl, nil),
		running:  prometheus.NewDesc("llmgw_endpoint_running", "Requests running on the engine, as last scraped.", lbl, nil),
		kv:       prometheus.NewDesc("llmgw_endpoint_kv_cache_usage", "Engine KV cache occupancy, 0..1.", lbl, nil),
		drain:    prometheus.NewDesc("llmgw_endpoint_drain_rate", "Observed completions per second.", lbl, nil),
		stale:    prometheus.NewDesc("llmgw_endpoint_metrics_stale", "1 when the endpoint's metrics are stale.", lbl, nil),
		breaker:  prometheus.NewDesc("llmgw_endpoint_breaker_open", "1 when the endpoint's circuit breaker is open.", lbl, nil),

		queueDepth: prometheus.NewDesc("llmgw_admission_queue_depth", "Requests waiting in each admission band.", []string{"band"}, nil),
		saturation: prometheus.NewDesc("llmgw_pool_saturation", "Pool saturation per model; 1 is the healthy limit.", []string{"model"}, nil),
	}
}

func (c *endpointCollector) Describe(ch chan<- *prometheus.Desc) {
	ch <- c.inflight
	ch <- c.waiting
	ch <- c.running
	ch <- c.kv
	ch <- c.drain
	ch <- c.stale
	ch <- c.breaker
	ch <- c.queueDepth
	ch <- c.saturation
}

func (c *endpointCollector) Collect(ch chan<- prometheus.Metric) {
	for _, s := range c.r.reg.Snapshots("") {
		pool := s.Pool
		if pool == "" {
			pool = "none"
		}
		g := func(d *prometheus.Desc, v float64) {
			ch <- prometheus.MustNewConstMetric(d, prometheus.GaugeValue, v, s.ID, pool)
		}
		g(c.inflight, float64(s.Inflight))
		g(c.waiting, float64(s.Stats.Waiting))
		g(c.running, float64(s.Stats.Running))
		g(c.kv, s.Stats.KVCacheUsage)
		g(c.drain, s.DrainRate)
		g(c.stale, boolToFloat(s.Stale))
		g(c.breaker, boolToFloat(s.Breaker.String() == "open"))
	}

	if c.r.adm != nil {
		for _, d := range c.r.adm.Depths() {
			ch <- prometheus.MustNewConstMetric(c.queueDepth, prometheus.GaugeValue, float64(d.Depth), d.Band)
		}
	}
	// Saturation is reported per model, which is the granularity admission
	// decisions are made at.
	for _, model := range c.r.servedModels() {
		if sat, ok := c.r.saturation(model, 0); ok {
			ch <- prometheus.MustNewConstMetric(c.saturation, prometheus.GaugeValue, sat, model)
		}
	}
}

func boolToFloat(b bool) float64 {
	if b {
		return 1
	}
	return 0
}

// MetricsHandler serves metrics and nothing else, for a listener that can be
// reached from outside the machine.
//
// The whole reason the admin listener is loopback is pprof: it can dump the
// process heap, and API keys are in that heap. Metrics are the one thing on
// that listener something else genuinely needs, so they get a door of their
// own rather than forcing a choice between being monitored and keeping the
// heap private.
//
// Returns nil before AdminHandler has been called, because that is where the
// registry is assembled — the same ordering that silently published nothing
// once already, which is why this says so rather than leaving a caller to find
// out from an empty page.
func (r *Router) MetricsHandler() http.Handler {
	if r.metricsHandler == nil {
		return nil
	}
	mux := http.NewServeMux()
	mux.Handle("GET /metrics", r.metricsHandler)
	// Health too, so a scrape target can be probed without a second port.
	mux.HandleFunc("GET /health", func(w http.ResponseWriter, _ *http.Request) {
		w.WriteHeader(http.StatusOK)
	})
	return mux
}