ITWorx-Pulse-Public/internal/livesampler/sampler.go

// Package livesampler binds the live WebSocket lane to the approved metric
// query path. It compiles an already planned queryplan.Request into one bounded
// instant Prometheus query and returns live samples whose freshness is
// classified by the same rules as every other metric response: missing or stale
// telemetry is reported explicitly and never as an empty successful result
// (ADR-0006, ADR-0008).
package livesampler

import (
	"context"
	"encoding/json"
	"errors"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/itworx/pulse/internal/live"
	"github.com/itworx/pulse/internal/metriccatalog"
	"github.com/itworx/pulse/internal/metricquery"
	"github.com/itworx/pulse/internal/prometheus"
	"github.com/itworx/pulse/internal/promqlbinding"
	"github.com/itworx/pulse/internal/queryplan"
)

// Reasons label the visible outcome of a sample that carries no fresh value.
const (
	ReasonNoData      = "no_data"
	ReasonSourceError = "source_error"
	ReasonTimeout     = "timeout"
	ReasonStale       = "stale"
)

const (
	reasonLabel        = "reason"
	defaultMaxSeries   = 20
	hardMaxSeries      = 500
	defaultTimeout     = 5 * time.Second
	hardMaxTimeout     = 30 * time.Second
	defaultStepSeconds = 60
	maxStepSeconds     = 86400
	pointsPerSeries    = 1
	maxLabels          = 20
	maxLabelLength     = 128
	maxSeriesName      = 256
	maxEpochSeconds    = 1e13
)

// Error is the stable, code-carrying rejection of a live sample request.
type Error struct {
	Code   string
	Field  string
	Detail string
}

func (e Error) Error() string {
	if e.Field == "" {
		return e.Code + ": " + e.Detail
	}
	return e.Code + " (" + e.Field + "): " + e.Detail
}

var (
	ErrSourceRequired  = Error{Code: "LIVE_SAMPLER_SOURCE_REQUIRED", Detail: "a metric source is required"}
	ErrCatalogRequired = Error{Code: "LIVE_SAMPLER_CATALOG_REQUIRED", Detail: "a loaded metric catalog is required"}
	ErrOptionsInvalid  = Error{Code: "LIVE_SAMPLER_OPTIONS_INVALID", Detail: "live sampler limits are out of bounds"}
)

// Source is the bounded instant-query surface of the Prometheus adapter.
type Source interface {
	Query(ctx context.Context, query string, at *time.Time) (prometheus.QueryResult, error)
}

// Options bound every upstream query the sampler is able to issue.
type Options struct {
	MaxSeries int
	Timeout   time.Duration
	Now       func() time.Time
}

// Sampler implements live.Sampler on top of the approved metric query path.
type Sampler struct {
	catalog   metriccatalog.Registry
	source    Source
	maxSeries int
	timeout   time.Duration
	now       func() time.Time
}

var _ live.Sampler = (*Sampler)(nil)

// New validates the sampler dependencies and limits up front so a misconfigured
// live lane fails at construction instead of at every sample.
func New(catalog metriccatalog.Registry, source Source, options Options) (*Sampler, error) {
	if source == nil {
		return nil, ErrSourceRequired
	}
	if len(catalog.Metrics()) == 0 {
		return nil, ErrCatalogRequired
	}
	if options.MaxSeries < 0 || options.MaxSeries > hardMaxSeries || options.Timeout < 0 || options.Timeout > hardMaxTimeout {
		return nil, ErrOptionsInvalid
	}
	maxSeries := options.MaxSeries
	if maxSeries == 0 {
		maxSeries = defaultMaxSeries
	}
	timeout := options.Timeout
	if timeout == 0 {
		timeout = defaultTimeout
	}
	now := options.Now
	if now == nil {
		now = func() time.Time { return time.Now().UTC() }
	}
	return &Sampler{catalog: catalog, source: source, maxSeries: maxSeries, timeout: timeout, now: now}, nil
}

// Sample executes one instant query for an already planned request. Upstream
// failures, timeouts, empty results and stale scrapes all return a visible
// sample carrying an explicit freshness and reason; an error is reserved for a
// request the sampler must refuse or a cancelled caller.
func (s *Sampler) Sample(ctx context.Context, request queryplan.Request) ([]live.Sample, error) {
	if s == nil || s.source == nil {
		return nil, ErrSourceRequired
	}
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	definition, ok := s.catalog.Find(strings.TrimSpace(request.Metric))
	if !ok {
		return nil, Error{Code: "LIVE_SAMPLE_METRIC_UNKNOWN", Field: "metric", Detail: "metric is not in the approved catalog"}
	}
	query, err := compile(definition, request)
	if err != nil {
		return nil, err
	}
	budget := s.seriesBudget(definition, request.MaxSeries)
	queryContext, cancel := context.WithTimeout(ctx, s.queryTimeout(definition))
	defer cancel()
	result, queryErr := s.source.Query(queryContext, query, nil)
	now := s.now()
	if queryErr != nil {
		if err := ctx.Err(); err != nil {
			return nil, err
		}
		return unavailable(definition, now, reasonFor(queryErr)), nil
	}
	samples, err := decode(definition, result.Data, now, budget)
	if err != nil {
		return nil, err
	}
	if len(samples) == 0 {
		return unavailable(definition, now, ReasonNoData), nil
	}
	return samples, nil
}

// compile reuses the approved template binding, so the sampler can only issue
// PromQL the metric catalog already sanctions for this metric and scope.
func compile(definition metriccatalog.Definition, request queryplan.Request) (string, error) {
	bounded := request
	bounded.Range.StepSeconds = stepSeconds(request)
	plan := queryplan.Plan{Metric: definition, Request: bounded, ResolvedScope: request.Scope}
	query, err := promqlbinding.CompilePlan(plan)
	if err != nil {
		return "", err
	}
	return query, nil
}

func stepSeconds(request queryplan.Request) int {
	step := request.Range.StepSeconds
	if step < 1 {
		return defaultStepSeconds
	}
	if step > maxStepSeconds {
		return maxStepSeconds
	}
	return step
}

// seriesBudget is the smallest of the sampler, metric and request budgets, so a
// subscription can narrow the cardinality of a live query but never widen it.
func (s *Sampler) seriesBudget(definition metriccatalog.Definition, requested int) int {
	budget := s.maxSeries
	if definition.Limits.MaxSeries > 0 && definition.Limits.MaxSeries < budget {
		budget = definition.Limits.MaxSeries
	}
	if definition.CardinalityBudget > 0 && definition.CardinalityBudget < budget {
		budget = definition.CardinalityBudget
	}
	if requested > 0 && requested < budget {
		budget = requested
	}
	if budget < 1 {
		return 1
	}
	return budget
}

func (s *Sampler) queryTimeout(definition metriccatalog.Definition) time.Duration {
	timeout := s.timeout
	metricTimeout := time.Duration(definition.Limits.TimeoutSeconds) * time.Second
	if metricTimeout > 0 && metricTimeout < timeout {
		timeout = metricTimeout
	}
	return timeout
}

func reasonFor(err error) string {
	if errors.Is(err, context.DeadlineExceeded) {
		return ReasonTimeout
	}
	var timeout interface{ Timeout() bool }
	if errors.As(err, &timeout) && timeout.Timeout() {
		return ReasonTimeout
	}
	return ReasonSourceError
}

func unavailable(definition metriccatalog.Definition, now time.Time, reason string) []live.Sample {
	return []live.Sample{{
		Series:    definition.SemanticName,
		Timestamp: now.UTC(),
		Freshness: metricquery.FreshnessUnknown,
		Labels:    map[string]string{reasonLabel: reason},
	}}
}

type payload struct {
	ResultType string          `json:"resultType"`
	Result     json.RawMessage `json:"result"`
}

type entry struct {
	Metric map[string]string   `json:"metric"`
	Value  []json.RawMessage   `json:"value"`
	Values [][]json.RawMessage `json:"values"`
}

// decode turns a bounded Prometheus instant response into at most one sample per
// series, capped by the series budget. A response above the budget is refused
// rather than silently truncated, like every other budget in the query path.
func decode(definition metriccatalog.Definition, data json.RawMessage, now time.Time, budget int) ([]live.Sample, error) {
	if len(data) == 0 {
		return nil, nil
	}
	var body payload
	if err := json.Unmarshal(data, &body); err != nil {
		return unavailable(definition, now, ReasonSourceError), nil
	}
	entries, err := entriesFor(body)
	if err != nil {
		return unavailable(definition, now, ReasonSourceError), nil
	}
	if len(entries) > budget {
		return nil, Error{Code: "LIVE_SAMPLE_SERIES_LIMIT", Field: "maxSeries", Detail: "live result exceeds the series budget"}
	}
	window := metricquery.FreshnessWindow(definition)
	samples := make([]live.Sample, 0, len(entries)*pointsPerSeries)
	for _, item := range entries {
		point, ok := latestPoint(item)
		if !ok {
			continue
		}
		observedAt, value, valid := parsePoint(point)
		if !valid {
			continue
		}
		labels := boundedLabels(item.Metric)
		sample := live.Sample{
			Series:    seriesName(definition.SemanticName, labels),
			Timestamp: observedAt.UTC(),
			Freshness: metricquery.ClassifyFreshness(observedAt, now, window),
			Labels:    labels,
		}
		if value != nil {
			sample.Value = value
		} else {
			sample.Freshness = metricquery.FreshnessUnknown
			sample.Labels[reasonLabel] = ReasonNoData
		}
		if sample.Freshness == metricquery.FreshnessStale {
			sample.Labels[reasonLabel] = ReasonStale
		}
		samples = append(samples, sample)
	}
	sort.Slice(samples, func(i, j int) bool { return samples[i].Series < samples[j].Series })
	return samples, nil
}

func entriesFor(body payload) ([]entry, error) {
	switch body.ResultType {
	case "vector", "matrix":
		var entries []entry
		if err := json.Unmarshal(body.Result, &entries); err != nil {
			return nil, err
		}
		return entries, nil
	case "scalar":
		var point []json.RawMessage
		if err := json.Unmarshal(body.Result, &point); err != nil {
			return nil, err
		}
		return []entry{{Value: point}}, nil
	default:
		return nil, errors.New("unsupported Prometheus result type")
	}
}

func latestPoint(item entry) ([]json.RawMessage, bool) {
	if len(item.Value) > 0 {
		return item.Value, true
	}
	if len(item.Values) > 0 {
		return item.Values[len(item.Values)-1], true
	}
	return nil, false
}

func parsePoint(point []json.RawMessage) (time.Time, *float64, bool) {
	if len(point) != 2 {
		return time.Time{}, nil, false
	}
	var seconds float64
	if err := json.Unmarshal(point[0], &seconds); err != nil {
		return time.Time{}, nil, false
	}
	if math.IsNaN(seconds) || math.IsInf(seconds, 0) || seconds < 0 || seconds > maxEpochSeconds {
		return time.Time{}, nil, false
	}
	whole := math.Floor(seconds)
	observedAt := time.Unix(int64(whole), int64(math.Round((seconds-whole)*float64(time.Second)))).UTC()
	var raw string
	if err := json.Unmarshal(point[1], &raw); err != nil {
		return observedAt, nil, true
	}
	value, err := strconv.ParseFloat(raw, 64)
	if err != nil || math.IsNaN(value) || math.IsInf(value, 0) {
		return observedAt, nil, true
	}
	return observedAt, &value, true
}

func boundedLabels(metric map[string]string) map[string]string {
	names := make([]string, 0, len(metric))
	for name := range metric {
		if name == "__name__" {
			continue
		}
		names = append(names, name)
	}
	sort.Strings(names)
	if len(names) > maxLabels {
		names = names[:maxLabels]
	}
	labels := make(map[string]string, len(names)+1)
	for _, name := range names {
		labels[bound(name, maxLabelLength)] = bound(metric[name], maxLabelLength)
	}
	return labels
}

func seriesName(semanticName string, labels map[string]string) string {
	if len(labels) == 0 {
		return semanticName
	}
	names := make([]string, 0, len(labels))
	for name := range labels {
		names = append(names, name)
	}
	sort.Strings(names)
	var builder strings.Builder
	builder.WriteString(semanticName)
	builder.WriteString("{")
	for index, name := range names {
		if index > 0 {
			builder.WriteString(",")
		}
		builder.WriteString(name)
		builder.WriteString("=")
		builder.WriteString(strconv.Quote(labels[name]))
	}
	builder.WriteString("}")
	return bound(builder.String(), maxSeriesName)
}

func bound(value string, limit int) string {
	if len(value) <= limit {
		return value
	}
	return value[:limit]
}