jarl/internal/stats/query.go

package stats

import (
	"cmp"
	"context"
	"fmt"
	"slices"
	"strings"
	"time"
)

// Querier answers questions about stored statistics.
type Querier struct {
	Store *Store
	Opts  Options
	Now   func() time.Time
}

// NewQuerier returns a querier for a collector's data.
func NewQuerier(c *Collector) *Querier {
	return &Querier{Store: c.store, Opts: c.opts, Now: time.Now}
}

func (q *Querier) loc() *time.Location {
	if q.Opts.Location != nil {
		return q.Opts.Location
	}
	return time.Local
}

// Resolution picks the finest stored resolution that covers a time range
// starting since ago.
func (q *Querier) Resolution(since time.Duration) string {
	switch {
	case since <= 48*time.Hour && since <= q.Opts.RetainMinute:
		return ResMinute
	case since <= q.Opts.RetainHour:
		return ResHour
	}
	return ResDay
}

// topsResolution picks the resolution for top lists, which exist per hour
// and per day.
func (q *Querier) topsResolution(since time.Duration) string {
	if since <= 7*24*time.Hour && since <= q.Opts.RetainHour {
		return ResHour
	}
	return ResDay
}

// bucketStart returns the start of the bucket containing t at a resolution.
func (q *Querier) bucketStart(res string, t time.Time) int64 {
	switch res {
	case ResMinute:
		return t.Unix() - t.Unix()%60
	case ResHour:
		return t.Unix() - t.Unix()%3600
	}
	l := t.In(q.loc())
	return time.Date(l.Year(), l.Month(), l.Day(), 0, 0, 0, 0, q.loc()).Unix()
}

// sums returns the totals of counters (optionally with a key prefix) in
// [from, to).
func (s *Store) sums(ctx context.Context, res string, from, to int64, prefix string) (map[string]int64, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT key, SUM(value) FROM counters WHERE res = ? AND ts >= ? AND ts < ? AND key LIKE ? ESCAPE '\' GROUP BY key`,
		res, from, to, likePrefix(prefix))
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[string]int64{}
	for rows.Next() {
		var k string
		var v int64
		if err := rows.Scan(&k, &v); err != nil {
			return nil, err
		}
		out[k] = v
	}
	return out, rows.Err()
}

func likePrefix(p string) string {
	r := strings.NewReplacer(`\`, `\\`, `%`, `\%`, `_`, `\_`)
	return r.Replace(p) + "%"
}

// series returns the values of a counter per bucket in [from, to).
func (s *Store) series(ctx context.Context, res string, from, to int64, key string) (map[int64]int64, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT ts, value FROM counters WHERE res = ? AND key = ? AND ts >= ? AND ts < ?`, res, key, from, to)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := map[int64]int64{}
	for rows.Next() {
		var ts, v int64
		if err := rows.Scan(&ts, &v); err != nil {
			return nil, err
		}
		out[ts] = v
	}
	return out, rows.Err()
}

// tops sums top lists in [from, to).
func (s *Store) tops(ctx context.Context, res, kind string, from, to int64, limit int) ([]Entry, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT key, SUM(count) AS n FROM tops WHERE res = ? AND kind = ? AND ts >= ? AND ts < ?
		GROUP BY key ORDER BY n DESC, key LIMIT ?`, res, kind, from, to, limit)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	out := []Entry{}
	for rows.Next() {
		var e Entry
		if err := rows.Scan(&e.Key, &e.Count); err != nil {
			return nil, err
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

func sorted(m map[string]int64, prefix string) []Entry {
	out := []Entry{}
	for k, v := range m {
		if strings.HasPrefix(k, prefix) && v != 0 {
			out = append(out, Entry{Key: strings.TrimPrefix(k, prefix), Count: v})
		}
	}
	slices.SortFunc(out, func(a, b Entry) int {
		if a.Count != b.Count {
			return cmp.Compare(b.Count, a.Count)
		}
		return cmp.Compare(a.Key, b.Key)
	})
	return out
}

// Summary is an overview of a time range.
type Summary struct {
	From       time.Time `json:"from"`
	To         time.Time `json:"to"`
	Resolution string    `json:"resolution"`
	Queries    int64     `json:"queries"`
	// AvgQPS is queries per second over the part of the range with data.
	AvgQPS float64 `json:"avg_qps"`
	Peak   *struct {
		Time    time.Time `json:"time"`
		Queries int64     `json:"queries"`
		Per     string    `json:"per"` // minute, hour or day
	} `json:"peak,omitempty"`
	RCodes        []Entry            `json:"rcodes"`
	QTypes        []Entry            `json:"qtypes"`
	Views         []Entry            `json:"views"`
	Zones         []Entry            `json:"zones"`
	Answers       []Entry            `json:"answers"`   // authoritative, recursive
	Transport     []Entry            `json:"transport"` // udp, tcp
	CacheHitRatio *float64           `json:"cache_hit_ratio,omitempty"`
	NXDomainRatio *float64           `json:"nxdomain_ratio,omitempty"`
	Tops          map[string][]Entry `json:"tops,omitempty"`
}

// Summary summarizes the last since.
func (q *Querier) Summary(ctx context.Context, since time.Duration, topLimit int) (*Summary, error) {
	now := q.Now()
	res := q.Resolution(since)
	from := q.bucketStart(res, now.Add(-since))
	to := now.Unix() + 1
	all, err := q.Store.sums(ctx, res, from, to, "")
	if err != nil {
		return nil, err
	}
	s := &Summary{From: time.Unix(from, 0), To: now, Resolution: res, Queries: all["queries"],
		RCodes: sorted(all, "rcode:"), QTypes: sorted(all, "qtype:"), Views: sorted(all, "view:"), Zones: sorted(all, "zone:"),
		Answers: sorted(all, "answer:"), Transport: sorted(all, "transport:")}
	if h, m := all["cache:hits"], all["cache:misses"]; h+m > 0 {
		r := float64(h) / float64(h+m)
		s.CacheHitRatio = &r
	}
	if s.Queries > 0 {
		r := float64(all["rcode:NXDOMAIN"]) / float64(s.Queries)
		s.NXDomainRatio = &r
	}
	per, err := q.Store.series(ctx, res, from, to, "queries")
	if err != nil {
		return nil, err
	}
	if len(per) > 0 {
		first, peakTS := int64(1<<62), int64(0)
		var peak int64 = -1
		for ts, v := range per {
			first = min(first, ts)
			if v > peak || v == peak && ts < peakTS {
				peak, peakTS = v, ts
			}
		}
		if secs := float64(to - first); secs > 0 {
			s.AvgQPS = float64(s.Queries) / secs
		}
		s.Peak = &struct {
			Time    time.Time `json:"time"`
			Queries int64     `json:"queries"`
			Per     string    `json:"per"`
		}{time.Unix(peakTS, 0), peak, map[string]string{ResMinute: "minute", ResHour: "hour", ResDay: "day"}[res]}
	}
	if topLimit > 0 {
		s.Tops = map[string][]Entry{}
		for _, k := range []string{KindName, KindClient, KindNXDomain} {
			es, err := q.Top(ctx, k, since, topLimit)
			if err != nil {
				return nil, err
			}
			s.Tops[k] = es
		}
	}
	return s, nil
}

// Top returns the largest entries of a top list over the last since.
func (q *Querier) Top(ctx context.Context, kind string, since time.Duration, limit int) ([]Entry, error) {
	if !slices.Contains(Kinds, kind) {
		return nil, fmt.Errorf("unknown list %q (use %s)", kind, strings.Join(Kinds, ", "))
	}
	now := q.Now()
	res := q.topsResolution(since)
	return q.Store.tops(ctx, res, kind, q.bucketStart(res, now.Add(-since)), now.Unix()+1, limit)
}

// Breakdown returns the totals of a counter group (qtype, rcode, view,
// zone, answer, transport, rtt, dnssec, cache, ratelimit) over the last since.
func (q *Querier) Breakdown(ctx context.Context, group string, since time.Duration) ([]Entry, error) {
	if !slices.Contains(Groups, group) {
		return nil, fmt.Errorf("unknown group %q (use %s)", group, strings.Join(Groups, ", "))
	}
	now := q.Now()
	res := q.Resolution(since)
	m, err := q.Store.sums(ctx, res, q.bucketStart(res, now.Add(-since)), now.Unix()+1, group+":")
	if err != nil {
		return nil, err
	}
	return sorted(m, group+":"), nil
}

// Groups are the counter groups Breakdown accepts.
var Groups = []string{"qtype", "rcode", "view", "zone", "answer", "transport", "rtt", "dnssec", "cache", "ratelimit"}

// Timeline is a set of series over time.
type Timeline struct {
	Step   string    `json:"step"`
	Times  []int64   `json:"times"` // bucket starts (Unix seconds)
	Series []*Series `json:"series"`
}

// Series is one counter over time.
type Series struct {
	Key    string  `json:"key"`
	Values []int64 `json:"values"`
}

var steps = []struct {
	name string
	d    time.Duration
	res  string
}{
	{"1m", time.Minute, ResMinute}, {"5m", 5 * time.Minute, ResMinute}, {"15m", 15 * time.Minute, ResMinute},
	{"1h", time.Hour, ResHour}, {"6h", 6 * time.Hour, ResHour}, {"1d", 24 * time.Hour, ResDay},
}

// Timeline returns counters over the last since. keys are counter names;
// a key ending in ":*" means every counter of that group. step is "auto"
// or one of 1m, 5m, 15m, 1h, 6h, 1d.
func (q *Querier) Timeline(ctx context.Context, since time.Duration, step string, keys []string) (*Timeline, error) {
	idx := -1
	for i, st := range steps {
		ok := st.res == ResMinute && since <= q.Opts.RetainMinute || st.res == ResHour && since <= q.Opts.RetainHour || st.res == ResDay
		if step == st.name {
			if !ok {
				return nil, fmt.Errorf("step %s needs %s data, which is only kept for %s", st.name, map[string]string{ResMinute: "per-minute", ResHour: "hourly"}[st.res], humanDuration(map[string]time.Duration{ResMinute: q.Opts.RetainMinute, ResHour: q.Opts.RetainHour}[st.res]))
			}
			idx = i
			break
		}
		if step == "auto" || step == "" {
			if ok && since/st.d <= 300 {
				idx = i
				break
			}
		}
	}
	if idx < 0 {
		if step != "auto" && step != "" {
			return nil, fmt.Errorf("unknown step %q (use auto, 1m, 5m, 15m, 1h, 6h or 1d)", step)
		}
		idx = len(steps) - 1
	}
	st := steps[idx]
	now := q.Now()
	// Buckets aligned to the step (days to local midnight).
	var starts []int64
	if st.res == ResDay {
		for t := time.Unix(q.bucketStart(ResDay, now.Add(-since)), 0).In(q.loc()); !t.After(now); t = t.AddDate(0, 0, 1) {
			starts = append(starts, t.Unix())
		}
	} else {
		sd := int64(st.d.Seconds())
		first := now.Add(-since).Unix()
		first -= first % sd
		for t := first; t <= now.Unix(); t += sd {
			starts = append(starts, t)
		}
	}
	from, to := starts[0], now.Unix()+1
	// Expand group keys.
	var expanded []string
	for _, k := range keys {
		if g, ok := strings.CutSuffix(k, ":*"); ok {
			m, err := q.Store.sums(ctx, st.res, from, to, g+":")
			if err != nil {
				return nil, err
			}
			for _, e := range sorted(m, "") {
				expanded = append(expanded, e.Key)
			}
			continue
		}
		expanded = append(expanded, k)
	}
	tl := &Timeline{Step: st.name, Times: starts, Series: []*Series{}}
	for _, k := range expanded {
		raw, err := q.Store.series(ctx, st.res, from, to, k)
		if err != nil {
			return nil, err
		}
		vals := make([]int64, len(starts))
		for ts, v := range raw {
			// The bucket is the last start <= ts.
			i, found := slices.BinarySearch(starts, ts)
			if !found {
				i--
			}
			if i >= 0 && i < len(vals) {
				vals[i] += v
			}
		}
		tl.Series = append(tl.Series, &Series{Key: k, Values: vals})
	}
	return tl, nil
}

// Heatmap is average activity per weekday and hour of day (local time).
type Heatmap struct {
	Key string `json:"key"`
	// Avg[d][h] is the average per hour on weekday d (0 = Monday) at
	// hour h; -1 where there is no data.
	Avg [7][24]float64 `json:"avg"`
	// Samples[d][h] is how many hours the average is based on.
	Samples [7][24]int `json:"samples"`
	From    time.Time  `json:"from"`
}

// Heatmap averages a counter (default "queries") by weekday and hour over
// the last since (at most the hourly retention).
func (q *Querier) Heatmap(ctx context.Context, key string, since time.Duration) (*Heatmap, error) {
	if key == "" {
		key = "queries"
	}
	since = min(since, q.Opts.RetainHour)
	now := q.Now()
	from := q.bucketStart(ResHour, now.Add(-since))
	raw, err := q.Store.series(ctx, ResHour, from, now.Unix()+1, key)
	if err != nil {
		return nil, err
	}
	hm := &Heatmap{Key: key, From: time.Unix(from, 0)}
	var sum [7][24]int64
	for ts, v := range raw {
		t := time.Unix(ts, 0).In(q.loc())
		d := (int(t.Weekday()) + 6) % 7
		sum[d][t.Hour()] += v
		hm.Samples[d][t.Hour()]++
	}
	for d := range 7 {
		for h := range 24 {
			if n := hm.Samples[d][h]; n > 0 {
				hm.Avg[d][h] = float64(sum[d][h]) / float64(n)
			} else {
				hm.Avg[d][h] = -1
			}
		}
	}
	return hm, nil
}

// ParseSince parses a range such as 90m, 24h, 7d or 2w.
func ParseSince(s string) (time.Duration, error) {
	if s == "" {
		return 24 * time.Hour, nil
	}
	mult := time.Duration(0)
	switch {
	case strings.HasSuffix(s, "d"):
		mult = 24 * time.Hour
	case strings.HasSuffix(s, "w"):
		mult = 7 * 24 * time.Hour
	}
	if mult > 0 {
		var n int
		if _, err := fmt.Sscanf(s[:len(s)-1], "%d", &n); err != nil || n <= 0 || fmt.Sprint(n) != s[:len(s)-1] {
			return 0, fmt.Errorf("invalid range %q (use e.g. 90m, 24h, 7d, 2w)", s)
		}
		return time.Duration(n) * mult, nil
	}
	d, err := time.ParseDuration(s)
	if err != nil || d <= 0 {
		return 0, fmt.Errorf("invalid range %q (use e.g. 90m, 24h, 7d, 2w)", s)
	}
	return d, nil
}