reforger_crawler_main/metrics/query.go

package metrics

import (
	"encoding/json"
	"fmt"
	"math"
	"sort"
	"strconv"
	"strings"
	"time"

	"gitea.tbdevent.eu/TBD/reforger_crawler_main/initializers"
	"gitea.tbdevent.eu/TBD/reforger_crawler_main/models"
)

const (
	maxQuerySeries = 50
	maxQueryPoints = 1500
)

// Selector picks series by service, name and exact label matches, and groups
// them by the GroupBy labels ("service" and "instance" work as labels too).
type Selector struct {
	Service string
	Name    string
	Match   map[string]string
	GroupBy []string
}

// Aggregations supported by Query and Totals.
var Aggregations = []string{"sum", "count", "rate", "avg", "min", "max", "last", "p50", "p90", "p95", "p99"}

func validAgg(agg string) bool {
	for _, a := range Aggregations {
		if a == agg {
			return true
		}
	}
	return false
}

// DefaultAgg picks a sensible aggregation for a metric kind.
func DefaultAgg(kind Kind) string {
	switch kind {
	case Counter:
		return "sum"
	case Gauge:
		return "last"
	case Histogram:
		return "p95"
	default:
		return "avg"
	}
}

// Point is [unix seconds, value or null] in JSON.
type Point struct {
	T int64
	V *float64
}

func (p Point) MarshalJSON() ([]byte, error) {
	if p.V == nil {
		return []byte("[" + strconv.FormatInt(p.T, 10) + ",null]"), nil
	}
	return []byte("[" + strconv.FormatInt(p.T, 10) + "," + strconv.FormatFloat(*p.V, 'g', 10, 64) + "]"), nil
}

type SeriesResult struct {
	Labels map[string]string `json:"labels"`
	Points []Point           `json:"points"`
}

type QueryResult struct {
	Name       string         `json:"name"`
	Kind       Kind           `json:"kind"`
	Agg        string         `json:"agg"`
	Resolution int64          `json:"resolution"`
	Step       int64          `json:"step"`
	From       int64          `json:"from"`
	To         int64          `json:"to"`
	Series     []SeriesResult `json:"series"`
	Truncated  bool           `json:"truncated"`
}

// matched is a series selected for a query, possibly only known from the DB.
type matched struct {
	id     uint
	live   *Series // nil when this process never saw the series
	labels map[string]string
	svc    string
	inst   string
	kind   Kind
	bounds []float64
}

func (m *matched) label(name string) string {
	switch name {
	case "service":
		if v, ok := m.labels["service"]; ok {
			return v
		}
		return m.svc
	case "instance":
		if v, ok := m.labels["instance"]; ok {
			return v
		}
		return m.inst
	}
	return m.labels[name]
}

func (sel Selector) matches(m *matched) bool {
	if sel.Service != "" && m.svc != sel.Service {
		return false
	}
	for k, v := range sel.Match {
		if m.label(k) != v {
			return false
		}
	}
	return true
}

func findSeries(sel Selector) ([]*matched, error) {
	byKey := map[seriesKey]*matched{}

	q := initializers.DB.Model(&models.MetricSeries{}).Where("name = ?", sel.Name)
	if sel.Service != "" {
		q = q.Where("service = ?", sel.Service)
	}
	var rows []models.MetricSeries
	if err := q.Find(&rows).Error; err != nil {
		return nil, err
	}
	for _, r := range rows {
		m := &matched{id: r.ID, svc: r.Service, inst: r.Instance, kind: Kind(r.Kind), labels: map[string]string{}}
		_ = json.Unmarshal([]byte(r.Labels), &m.labels)
		if r.Bounds != "" {
			_ = json.Unmarshal([]byte(r.Bounds), &m.bounds)
		}
		byKey[seriesKey{r.Service, r.Instance, r.Name, r.Labels}] = m
	}

	for _, s := range reg.all() {
		if s.Name != sel.Name || (sel.Service != "" && s.Service != sel.Service) {
			continue
		}
		k := s.key()
		if m := byKey[k]; m != nil {
			m.live = s
			continue
		}
		byKey[k] = &matched{live: s, svc: s.Service, inst: s.Instance, kind: s.Kind, bounds: s.Bounds, labels: s.Labels}
	}

	out := make([]*matched, 0, len(byKey))
	for _, m := range byKey {
		if sel.matches(m) {
			out = append(out, m)
		}
	}
	return out, nil
}

type groupAcc struct {
	labels  map[string]string
	steps   map[int64]*Agg
	lastSum map[int64]float64
}

// collect loads all matching data in [from, to] and folds it into groups and
// step buckets chosen by bucketFn. Gauge "last" values are summed across the
// series of a group (e.g. queue depth over all instances).
func collect(sel Selector, from, to, res int64, bucketFn func(int64) int64) (map[string]*groupAcc, Kind, []float64, error) {
	series, err := findSeries(sel)
	if err != nil {
		return nil, "", nil, err
	}
	if len(series) == 0 {
		return map[string]*groupAcc{}, "", nil, nil
	}
	kind, bounds := series[0].kind, series[0].bounds

	perSeries := map[*matched]map[int64]*Agg{}
	add := func(m *matched, ts int64, a *Agg) {
		b := bucketFn(ts)
		steps := perSeries[m]
		if steps == nil {
			steps = map[int64]*Agg{}
			perSeries[m] = steps
		}
		if cur := steps[b]; cur != nil {
			cur.merge(a)
		} else {
			steps[b] = a.clone()
		}
	}

	byID := map[uint]*matched{}
	var ids []uint
	for _, m := range series {
		if m.id != 0 {
			byID[m.id] = m
			ids = append(ids, m.id)
		}
	}

	load := func(resolution int, lo, hi int64) error {
		if lo > hi || len(ids) == 0 {
			return nil
		}
		for start := 0; start < len(ids); start += 300 {
			end := min(start+300, len(ids))
			var rows []models.MetricPoint
			err := initializers.DB.
				Where("resolution = ? AND series_id IN ? AND bucket >= ? AND bucket <= ?", resolution, ids[start:end], lo, hi).
				Find(&rows).Error
			if err != nil {
				return err
			}
			for i := range rows {
				add(byID[rows[i].SeriesID], rows[i].Bucket, pointToAgg(&rows[i]))
			}
		}
		return nil
	}

	if res == resHour {
		cut := rolledUntil.Load()
		if err := load(resHour, hourOf(from), cut-1); err != nil {
			return nil, "", nil, err
		}
		if err := load(resMinute, max(cut, from), to); err != nil {
			return nil, "", nil, err
		}
	} else if err := load(resMinute, minuteOf(from), to); err != nil {
		return nil, "", nil, err
	}

	// not yet flushed buckets
	live := map[*Series]*matched{}
	want := map[*Series]bool{}
	for _, m := range series {
		if m.live != nil {
			live[m.live] = m
			want[m.live] = true
		}
	}
	for k, a := range defaultAgg.snapshotOpen(want) {
		if k.bucket >= minuteOf(from) && k.bucket <= to {
			add(live[k.s], k.bucket, a)
		}
	}

	groups := map[string]*groupAcc{}
	for m, steps := range perSeries {
		labels := make(map[string]string, len(sel.GroupBy))
		parts := make([]string, len(sel.GroupBy))
		for i, g := range sel.GroupBy {
			labels[g] = m.label(g)
			parts[i] = labels[g]
		}
		gk := strings.Join(parts, "\x00")
		g := groups[gk]
		if g == nil {
			g = &groupAcc{labels: labels, steps: map[int64]*Agg{}, lastSum: map[int64]float64{}}
			groups[gk] = g
		}
		for b, a := range steps {
			if cur := g.steps[b]; cur != nil {
				cur.merge(a)
			} else {
				g.steps[b] = a.clone()
			}
			g.lastSum[b] += a.Last
		}
	}
	return groups, kind, bounds, nil
}

func evaluate(a *Agg, lastSum float64, agg string, kind Kind, bounds []float64, span int64) (float64, bool) {
	if a == nil || a.Count == 0 {
		return 0, false
	}
	switch agg {
	case "count":
		return float64(a.Count), true
	case "sum":
		return a.Sum, true
	case "rate":
		if span <= 0 {
			return 0, false
		}
		if kind == Counter {
			return a.Sum / float64(span), true
		}
		return float64(a.Count) / float64(span), true
	case "avg":
		return a.Sum / float64(a.Count), true
	case "min":
		return a.Min, true
	case "max":
		return a.Max, true
	case "last":
		return lastSum, true
	case "p50", "p90", "p95", "p99":
		q, _ := strconv.ParseFloat(agg[1:], 64)
		return quantile(a, bounds, q/100)
	}
	return 0, false
}

func pickResolution(from, to, step int64) int64 {
	minuteHorizon := time.Now().Unix() - int64(initializers.MetricsMinuteRetentionDays)*86400
	if to-from > 48*3600 || step >= resHour || from < minuteHorizon {
		return resHour
	}
	return resMinute
}

// Query returns one time series per group with values at every step.
func Query(sel Selector, agg string, from, to, step int64) (*QueryResult, error) {
	if sel.Name == "" {
		return nil, fmt.Errorf("name is required")
	}
	if to <= from {
		return nil, fmt.Errorf("to must be after from")
	}
	res := pickResolution(from, to, step)
	if step < res {
		step = res
	}
	if (to-from)/step > maxQueryPoints {
		step = (to - from) / maxQueryPoints
	}
	step = int64(math.Ceil(float64(step)/float64(res))) * res
	align := func(ts int64) int64 { return ts - ((ts%step)+step)%step }

	groups, kind, bounds, err := collect(sel, from, to, res, align)
	if err != nil {
		return nil, err
	}
	if agg == "" {
		agg = DefaultAgg(kind)
	}
	if !validAgg(agg) {
		return nil, fmt.Errorf("unknown aggregation %q", agg)
	}

	result := &QueryResult{Name: sel.Name, Kind: kind, Agg: agg, Resolution: res, Step: step, From: from, To: to, Series: []SeriesResult{}}
	keys := make([]string, 0, len(groups))
	for k := range groups {
		keys = append(keys, k)
	}
	sort.Strings(keys)
	if len(keys) > maxQuerySeries {
		keys = keys[:maxQuerySeries]
		result.Truncated = true
	}
	for _, k := range keys {
		g := groups[k]
		sr := SeriesResult{Labels: g.labels}
		for b := align(from); b <= to; b += step {
			p := Point{T: b}
			if v, ok := evaluate(g.steps[b], g.lastSum[b], agg, kind, bounds, step); ok && !math.IsNaN(v) && !math.IsInf(v, 0) {
				p.V = &v
			}
			sr.Points = append(sr.Points, p)
		}
		result.Series = append(result.Series, sr)
	}
	return result, nil
}

// Total is the whole-range aggregate of one group.
type Total struct {
	Labels map[string]string `json:"labels"`
	Count  int64             `json:"count"`
	Sum    float64           `json:"sum"`
	Min    *float64          `json:"min"`
	Max    *float64          `json:"max"`
	Avg    *float64          `json:"avg"`
	Last   *float64          `json:"last"`
	Rate   *float64          `json:"rate"`
	P50    *float64          `json:"p50"`
	P95    *float64          `json:"p95"`
	P99    *float64          `json:"p99"`
}

// Totals aggregates [from, to] into one value set per group.
func Totals(sel Selector, from, to int64) ([]Total, error) {
	res := pickResolution(from, to, 0)
	groups, kind, bounds, err := collect(sel, from, to, res, func(int64) int64 { return 0 })
	if err != nil {
		return nil, err
	}
	span := to - from
	out := make([]Total, 0, len(groups))
	for _, g := range groups {
		a := g.steps[0]
		if a == nil {
			continue
		}
		t := Total{Labels: g.labels, Count: a.Count, Sum: a.Sum}
		val := func(agg string) *float64 {
			if v, ok := evaluate(a, g.lastSum[0], agg, kind, bounds, span); ok && !math.IsNaN(v) && !math.IsInf(v, 0) {
				return &v
			}
			return nil
		}
		t.Min, t.Max, t.Avg, t.Last, t.Rate = val("min"), val("max"), val("avg"), val("last"), val("rate")
		t.P50, t.P95, t.P99 = val("p50"), val("p95"), val("p99")
		out = append(out, t)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Count > out[j].Count })
	return out, nil
}

// Sum returns the summed value of a counter (or count of observations for
// other kinds) over the range, across all matching series.
func Sum(sel Selector, from, to int64) float64 {
	sel.GroupBy = nil
	totals, err := Totals(sel, from, to)
	if err != nil || len(totals) == 0 {
		return 0
	}
	if f := familyByName(sel.Name); f != nil && f.Kind != Counter {
		return float64(totals[0].Count)
	}
	return totals[0].Sum
}

// CatalogEntry describes one metric name of one service.
type CatalogEntry struct {
	Service   string              `json:"service"`
	Name      string              `json:"name"`
	Kind      Kind                `json:"kind"`
	Help      string              `json:"help"`
	Labels    map[string][]string `json:"labels"`
	Instances []string            `json:"instances"`
	Series    int                 `json:"series"`
	LastSeen  int64               `json:"lastSeen"`
}

// Catalog lists every metric known in the DB or in memory.
func Catalog(service string) ([]CatalogEntry, error) {
	type acc struct {
		e      CatalogEntry
		labels map[string]map[string]bool
		inst   map[string]bool
		keys   map[seriesKey]bool
	}
	entries := map[string]*acc{}
	addSeries := func(svc, inst, name string, kind Kind, help, labelsJSON string, labels map[string]string, lastSeen int64) {
		k := svc + "/" + name
		a := entries[k]
		if a == nil {
			a = &acc{e: CatalogEntry{Service: svc, Name: name, Kind: kind, Help: help}, labels: map[string]map[string]bool{}, inst: map[string]bool{}, keys: map[seriesKey]bool{}}
			entries[k] = a
		}
		if a.e.Help == "" {
			a.e.Help = help
		}
		sk := seriesKey{svc, inst, name, labelsJSON}
		if !a.keys[sk] {
			a.keys[sk] = true
			a.e.Series++
		}
		a.inst[inst] = true
		for lk, lv := range labels {
			if a.labels[lk] == nil {
				a.labels[lk] = map[string]bool{}
			}
			a.labels[lk][lv] = true
		}
		if lastSeen > a.e.LastSeen {
			a.e.LastSeen = lastSeen
		}
	}

	q := initializers.DB.Model(&models.MetricSeries{})
	if service != "" {
		q = q.Where("service = ?", service)
	}
	var rows []models.MetricSeries
	if err := q.Find(&rows).Error; err != nil {
		return nil, err
	}
	for _, r := range rows {
		labels := map[string]string{}
		_ = json.Unmarshal([]byte(r.Labels), &labels)
		addSeries(r.Service, r.Instance, r.Name, Kind(r.Kind), r.Help, r.Labels, labels, r.LastSeen)
	}
	for _, s := range reg.all() {
		if service != "" && s.Service != service {
			continue
		}
		addSeries(s.Service, s.Instance, s.Name, s.Kind, s.Help, s.LabelsJSON, s.Labels, reg.seen(s))
	}

	out := make([]CatalogEntry, 0, len(entries))
	for _, a := range entries {
		a.e.Labels = map[string][]string{}
		for lk, vals := range a.labels {
			list := make([]string, 0, len(vals))
			for v := range vals {
				list = append(list, v)
			}
			sort.Strings(list)
			a.e.Labels[lk] = list
		}
		for inst := range a.inst {
			a.e.Instances = append(a.e.Instances, inst)
		}
		sort.Strings(a.e.Instances)
		out = append(out, a.e)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Service != out[j].Service {
			return out[i].Service < out[j].Service
		}
		return out[i].Name < out[j].Name
	})
	return out, nil
}