internal/agg/agg.go - master - rustwizard/pgtrace

internal/agg/agg.go
206 строк4 KB
Rust Wizard
feat(offcpu): aggregate wait types per query and expose in console/agent
08 авг 2026, 09:51
08 авг 2026, 09:5173d29f7
// Package agg aggregates syscall latency events into fixed time windows.
package agg

import (
	"maps"
	"sort"
	"strings"
	"sync"
	"time"

	"github.com/rustwizard/pgtrace/internal/bpf"
)

const (
	maxQueryLen = 120
	noQuery     = "?"
)

// Stats holds aggregated syscall observations within one window.
type Stats struct {
	Count   int
	TotalMs float64
	MaxMs   float64
	// Queries maps normalized query text to call count; noQuery ("?")
	// buckets events with no known query.
	Queries map[string]int
}

// Add records one observation of duration d attributed to query.
func (s *Stats) Add(d time.Duration, query string) {
	ms := float64(d) / float64(time.Millisecond)
	s.Count++

	s.TotalMs += ms
	if ms > s.MaxMs {
		s.MaxMs = ms
	}

	q := normalizeQuery(query)
	if q == "" {
		q = noQuery
	}

	s.Queries[q]++
}

// TopQuery returns the most frequent query and its count.
// Returns ("", 0) if no events were recorded.
func (s *Stats) TopQuery() (string, int) {
	top, n := "", 0
	for q, c := range s.Queries {
		if c > n {
			top, n = q, c
		}
	}

	return top, n
}

// QueryEntry is one query with its call count, ordered most frequent first.
type QueryEntry struct {
	Query string
	Count int
}

// WaitStat holds aggregated off-CPU time for one wait type within a window.
type WaitStat struct {
	TotalMs float64
	// Queries maps normalized query text to the time spent waiting in this
	// wait type while running that query.
	Queries map[string]float64
}

// TopQueries returns the n most frequent queries, ordered by count descending.
func (s *Stats) TopQueries(n int) []QueryEntry {
	if n < 1 {
		return nil
	}

	entries := make([]QueryEntry, 0, len(s.Queries))
	for q, c := range s.Queries {
		entries = append(entries, QueryEntry{Query: q, Count: c})
	}

	sort.Slice(entries, func(i, j int) bool {
		if entries[i].Count != entries[j].Count {
			return entries[i].Count > entries[j].Count
		}

		return entries[i].Query < entries[j].Query
	})

	if len(entries) > n {
		entries = entries[:n]
	}

	return entries
}

// Aggregator accumulates events per syscall for the current window.
// Safe for concurrent use by the event loop and the window ticker.
type Aggregator struct {
	mu    sync.Mutex
	start time.Time
	m     map[uint32]*Stats
	w     map[string]*WaitStat
}

// New returns an Aggregator with the window starting at now.
func New(now time.Time) *Aggregator {
	return &Aggregator{start: now, m: make(map[uint32]*Stats), w: make(map[string]*WaitStat)}
}

// Add records ev attributed to query (may be empty).
func (a *Aggregator) Add(ev bpf.Event, query string) {
	a.mu.Lock()
	defer a.mu.Unlock()

	s, ok := a.m[ev.SyscallID]
	if !ok {
		s = &Stats{Queries: make(map[string]int)}
		a.m[ev.SyscallID] = s
	}

	//nolint:gosec // kernel duration always fits in int64 nanoseconds
	s.Add(time.Duration(ev.DurationNs), query)
}

// AddWait records off-CPU time of waitType attributed to query (may be empty).
func (a *Aggregator) AddWait(waitType string, d time.Duration, query string) {
	a.mu.Lock()
	defer a.mu.Unlock()

	ws, ok := a.w[waitType]
	if !ok {
		ws = &WaitStat{Queries: make(map[string]float64)}
		a.w[waitType] = ws
	}

	ms := float64(d) / float64(time.Millisecond)
	ws.TotalMs += ms

	q := normalizeQuery(query)
	if q == "" {
		q = noQuery
	}

	ws.Queries[q] += ms
}

// Snapshot returns the window bounds and a copy of per-syscall stats plus
// per-wait-type stats, then resets the window to end at end. Stats maps are
// nil when empty.
func (a *Aggregator) Snapshot(end time.Time) (time.Time, time.Time, map[uint32]Stats, map[string]WaitStat) {
	a.mu.Lock()
	defer a.mu.Unlock()

	start := a.start

	a.start = end

	var syscalls map[uint32]Stats
	if len(a.m) > 0 {
		syscalls = make(map[uint32]Stats, len(a.m))
		for id, s := range a.m {
			cp := Stats{
				Count:   s.Count,
				TotalMs: s.TotalMs,
				MaxMs:   s.MaxMs,
				Queries: make(map[string]int, len(s.Queries)),
			}
			maps.Copy(cp.Queries, s.Queries)

			syscalls[id] = cp
		}
	}

	clear(a.m)

	var waits map[string]WaitStat
	if len(a.w) > 0 {
		waits = make(map[string]WaitStat, len(a.w))
		for wt, ws := range a.w {
			cp := WaitStat{
				TotalMs: ws.TotalMs,
				Queries: make(map[string]float64, len(ws.Queries)),
			}
			maps.Copy(cp.Queries, ws.Queries)

			waits[wt] = cp
		}
	}

	clear(a.w)

	return start, end, syscalls, waits
}

func normalizeQuery(q string) string {
	q = strings.Join(strings.Fields(q), " ")
	if len(q) > maxQueryLen {
		q = q[:maxQueryLen] + "..."
	}

	return q
}