llmgw/internal/router/scrape/poller.go

package scrape

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"net/http"
	"sync"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/router/endpoint"
)

// Config holds poller settings.
type Config struct {
	// Interval is the scrape period. Plan §7.3 suggests 1-2s.
	Interval time.Duration
	// Timeout bounds a single scrape; it must be shorter than Interval.
	Timeout time.Duration
	// Path is the metrics endpoint.
	Path string
	// Mapping selects the metric names for endpoints that name no engine. An
	// endpoint that names one is read with that engine's mapping instead.
	Mapping Mapping
}

// DefaultConfig returns sensible defaults.
func DefaultConfig() Config {
	return Config{
		Interval: 2 * time.Second,
		Timeout:  time.Second,
		Path:     "/metrics",
		Mapping:  VLLMMapping(),
	}
}

func (c *Config) normalize() {
	if c.Interval <= 0 {
		c.Interval = 2 * time.Second
	}
	if c.Timeout <= 0 || c.Timeout >= c.Interval {
		c.Timeout = c.Interval / 2
	}
	if c.Path == "" {
		c.Path = "/metrics"
	}
	c.Mapping.normalize()
}

// Poller scrapes each registered endpoint periodically, one goroutine each.
//
// Per-endpoint goroutines plus jitter: scraping every endpoint at the same
// instant loads both the engines and the router in bursts and synchronizes the
// measurements, which makes the scores oscillate in lockstep.
type Poller struct {
	cfg  Config
	reg  *endpoint.Registry
	http *http.Client
	log  *slog.Logger

	mu      sync.Mutex
	running map[string]context.CancelFunc
	wg      sync.WaitGroup
}

// NewPoller creates a poller. Call Sync to start it.
func NewPoller(reg *endpoint.Registry, cfg Config, log *slog.Logger) *Poller {
	cfg.normalize()
	if log == nil {
		log = slog.Default()
	}
	return &Poller{
		cfg: cfg,
		reg: reg,
		log: log,
		http: &http.Client{
			Timeout: cfg.Timeout,
			Transport: &http.Transport{
				MaxIdleConnsPerHost: 2,
				DisableCompression:  true,
			},
		},
		running: make(map[string]context.CancelFunc),
	}
}

// Sync starts and stops scrape goroutines to match the registry's endpoint
// list. Call it again whenever that list changes.
func (p *Poller) Sync(ctx context.Context) {
	infos := p.reg.Infos()

	p.mu.Lock()
	defer p.mu.Unlock()

	want := make(map[string]endpoint.Info, len(infos))
	for _, in := range infos {
		if in.Engine == endpoint.EngineExternal {
			continue // a provider elsewhere publishes no /metrics
		}
		want[in.ID] = in
	}

	for id, cancel := range p.running {
		if _, ok := want[id]; !ok {
			cancel()
			delete(p.running, id)
		}
	}
	for id, in := range want {
		if _, ok := p.running[id]; ok {
			continue
		}
		cctx, cancel := context.WithCancel(ctx)
		p.running[id] = cancel
		p.wg.Add(1)
		go func() {
			defer p.wg.Done()
			p.loop(cctx, in)
		}()
	}
}

// Stop halts every scrape goroutine and waits for them to finish.
func (p *Poller) Stop() {
	p.mu.Lock()
	for id, cancel := range p.running {
		cancel()
		delete(p.running, id)
	}
	p.mu.Unlock()
	p.wg.Wait()
}

func (p *Poller) loop(ctx context.Context, in endpoint.Info) {
	// Jitter keeps goroutines that start together from staying in lockstep.
	jitter := time.Duration(rand.Int64N(int64(p.cfg.Interval))) //nolint:gosec // jitter
	select {
	case <-ctx.Done():
		return
	case <-time.After(jitter):
	}

	t := time.NewTicker(p.cfg.Interval)
	defer t.Stop()

	// One goroutine per endpoint, so this needs no lock.
	var reported string

	p.once(ctx, in, &reported)
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			p.once(ctx, in, &reported)
		}
	}
}

func (p *Poller) once(ctx context.Context, in endpoint.Info, reported *string) {
	stats, err := p.fetch(ctx, in.URL, in.APIKey, p.mappingFor(in))
	if err != nil {
		if ctx.Err() != nil {
			return
		}
		p.reg.ScrapeFailed(in.ID)
		p.report(in, err, reported)
		return
	}
	if *reported != "" {
		p.log.Info("endpoint metrics are readable again", "endpoint", in.ID)
		*reported = ""
	}
	p.reg.Observe(in.ID, stats)
}

// report logs a scrape failure the first time it appears and then stays quiet
// until the reason changes.
//
// Per-scrape debug logging was not enough. A contract violation is a
// compatibility problem, not a blip: every scrape fails, the router has no
// numbers for that endpoint and schedules to it blind, and at the default log
// level nothing says so. A cluster can look entirely healthy while none of its
// routing decisions are informed by anything.
func (p *Poller) report(in endpoint.Info, err error, reported *string) {
	var missing *MissingError
	if errors.As(err, &missing) {
		msg := err.Error()
		if *reported == msg {
			return
		}
		*reported = msg
		p.log.Warn("engine does not serve the metrics this endpoint is read with; "+
			"routing to it is uninformed",
			"endpoint", in.ID, "engine", in.Engine, "missing", missing.Missing)
		return
	}

	// Ordinary failures — a restarting engine, a timeout — are expected and
	// transient, so they stay quiet unless someone is looking.
	p.log.Debug("scrape failed", "endpoint", in.ID, "err", err)
	*reported = "unreachable"
}

func (p *Poller) fetch(ctx context.Context, base, apiKey string, mapping Mapping) (endpoint.Stats, error) {
	ctx, cancel := context.WithTimeout(ctx, p.cfg.Timeout)
	defer cancel()

	path := p.cfg.Path
	if mapping.Path != "" {
		path = mapping.Path
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, base+path, nil)
	if err != nil {
		return endpoint.Stats{}, err
	}
	// llama.cpp guards /metrics with its key, and llmgw's TensorRT-LLM
	// launcher guards everything but /health; vLLM does not, and ignores one.
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}
	resp, err := p.http.Do(req)
	if err != nil {
		return endpoint.Stats{}, err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return endpoint.Stats{}, fmt.Errorf("unexpected status: %s", resp.Status)
	}
	return Parse(resp.Body, mapping)
}

// mappingFor picks the vocabulary this endpoint's metrics are written in.
//
// The configured mapping wins for endpoints that name no engine, so an operator
// can still override the names when a version renames them (plan §14).
func (p *Poller) mappingFor(in endpoint.Info) Mapping {
	if in.Engine == "" {
		return p.cfg.Mapping
	}
	return MappingFor(in.Engine)
}