reforger_crawler_main/metrics/aggregator.go

package metrics

import (
	"encoding/json"
	"math"
	"sort"
	"sync"
	"sync/atomic"
)

// disabled turns all recording into no-ops (METRICS_ENABLED=false).
var disabled atomic.Bool

// SetEnabled switches metric recording on or off.
func SetEnabled(on bool) { disabled.Store(!on) }

// Agg is the aggregate of all samples of one series inside one bucket.
type Agg struct {
	Count  int64
	Sum    float64
	Min    float64
	Max    float64
	Last   float64
	LastTs int64
	Hist   []uint64 // len(bounds)+1, last element is +Inf
}

func (a *Agg) add(v float64, ts int64, bounds []float64) {
	if a.Count == 0 {
		a.Min, a.Max = v, v
	} else {
		a.Min = math.Min(a.Min, v)
		a.Max = math.Max(a.Max, v)
	}
	a.Count++
	a.Sum += v
	if ts >= a.LastTs {
		a.Last, a.LastTs = v, ts
	}
	if len(bounds) > 0 {
		if a.Hist == nil {
			a.Hist = make([]uint64, len(bounds)+1)
		}
		a.Hist[bucketIndex(bounds, v)]++
	}
}

func (a *Agg) addSummary(count int64, sum, min, max float64, ts int64) {
	if count <= 0 {
		return
	}
	if a.Count == 0 {
		a.Min, a.Max = min, max
	} else {
		a.Min = math.Min(a.Min, min)
		a.Max = math.Max(a.Max, max)
	}
	a.Count += count
	a.Sum += sum
	if ts >= a.LastTs {
		a.Last, a.LastTs = sum/float64(count), ts
	}
}

// merge folds b into a.
func (a *Agg) merge(b *Agg) {
	if b == nil || b.Count == 0 {
		return
	}
	if a.Count == 0 {
		a.Min, a.Max = b.Min, b.Max
	} else {
		a.Min = math.Min(a.Min, b.Min)
		a.Max = math.Max(a.Max, b.Max)
	}
	a.Count += b.Count
	a.Sum += b.Sum
	if b.LastTs >= a.LastTs {
		a.Last, a.LastTs = b.Last, b.LastTs
	}
	if len(b.Hist) > 0 {
		if len(a.Hist) == 0 {
			a.Hist = make([]uint64, len(b.Hist))
		}
		if len(a.Hist) == len(b.Hist) {
			for i, v := range b.Hist {
				a.Hist[i] += v
			}
		}
	}
}

func (a *Agg) clone() *Agg {
	c := *a
	if a.Hist != nil {
		c.Hist = append([]uint64(nil), a.Hist...)
	}
	return &c
}

func (a *Agg) histJSON() string {
	if len(a.Hist) == 0 {
		return ""
	}
	b, _ := json.Marshal(a.Hist)
	return string(b)
}

func bucketIndex(bounds []float64, v float64) int {
	return sort.SearchFloat64s(bounds, v) // first bound >= v; len(bounds) = +Inf
}

// quantile estimates the q-quantile from histogram buckets with linear
// interpolation inside the bucket (like Prometheus), clamped to [min, max].
func quantile(a *Agg, bounds []float64, q float64) (float64, bool) {
	if a == nil || a.Count == 0 || len(a.Hist) != len(bounds)+1 {
		return 0, false
	}
	var total uint64
	for _, c := range a.Hist {
		total += c
	}
	if total == 0 {
		return 0, false
	}
	rank := q * float64(total)
	var cum uint64
	for i, c := range a.Hist {
		if c == 0 {
			continue
		}
		if float64(cum+c) >= rank {
			lower := a.Min
			if i > 0 {
				lower = math.Max(bounds[i-1], a.Min)
			}
			upper := a.Max
			if i < len(bounds) {
				upper = math.Min(bounds[i], a.Max)
			}
			if upper < lower {
				upper = lower
			}
			frac := (rank - float64(cum)) / float64(c)
			return lower + (upper-lower)*frac, true
		}
		cum += c
	}
	return a.Max, true
}

type pointKey struct {
	s      *Series
	bucket int64
}

// cumulative is the since-start state of a series, used for /metrics.
type cumulative struct {
	count int64
	sum   float64
	last  float64
	hist  []uint64
}

// aggregator holds the not yet persisted minute buckets and cumulative state.
type aggregator struct {
	mu   sync.Mutex
	open map[pointKey]*Agg
	cum  map[*Series]*cumulative
}

var defaultAgg = &aggregator{
	open: make(map[pointKey]*Agg),
	cum:  make(map[*Series]*cumulative),
}

func minuteOf(ts int64) int64 { return ts - ts%60 }

func (g *aggregator) entry(s *Series, ts int64) (*Agg, *cumulative) {
	k := pointKey{s, minuteOf(ts)}
	a := g.open[k]
	if a == nil {
		a = &Agg{}
		g.open[k] = a
	}
	c := g.cum[s]
	if c == nil {
		c = &cumulative{}
		g.cum[s] = c
	}
	return a, c
}

func (g *aggregator) observe(s *Series, ts int64, v float64) {
	if disabled.Load() {
		return
	}
	g.mu.Lock()
	a, c := g.entry(s, ts)
	a.add(v, ts, s.Bounds)
	c.count++
	c.sum += v
	c.last = v
	if len(s.Bounds) > 0 {
		if c.hist == nil {
			c.hist = make([]uint64, len(s.Bounds)+1)
		}
		c.hist[bucketIndex(s.Bounds, v)]++
	}
	g.mu.Unlock()
	reg.touch(s, ts)
}

func (g *aggregator) observeSummary(s *Series, ts int64, count int64, sum, min, max float64) {
	if disabled.Load() {
		return
	}
	g.mu.Lock()
	a, c := g.entry(s, ts)
	a.addSummary(count, sum, min, max, ts)
	c.count += count
	c.sum += sum
	c.last = sum / float64(count)
	g.mu.Unlock()
	reg.touch(s, ts)
}

// takeClosed removes and returns all buckets that started before `before`.
func (g *aggregator) takeClosed(before int64) map[pointKey]*Agg {
	g.mu.Lock()
	defer g.mu.Unlock()
	out := make(map[pointKey]*Agg)
	for k, a := range g.open {
		if k.bucket < before {
			out[k] = a
			delete(g.open, k)
		}
	}
	return out
}

// putBack re-queues buckets after a failed flush.
func (g *aggregator) putBack(points map[pointKey]*Agg) {
	g.mu.Lock()
	defer g.mu.Unlock()
	for k, a := range points {
		if cur := g.open[k]; cur != nil {
			cur.merge(a)
		} else {
			g.open[k] = a
		}
	}
}

// snapshotOpen copies the open buckets of the given series.
func (g *aggregator) snapshotOpen(want map[*Series]bool) map[pointKey]*Agg {
	g.mu.Lock()
	defer g.mu.Unlock()
	out := make(map[pointKey]*Agg)
	for k, a := range g.open {
		if want[k.s] {
			out[k] = a.clone()
		}
	}
	return out
}

func (g *aggregator) openCount() int {
	g.mu.Lock()
	defer g.mu.Unlock()
	return len(g.open)
}

func (g *aggregator) snapshotCum() map[*Series]cumulative {
	g.mu.Lock()
	defer g.mu.Unlock()
	out := make(map[*Series]cumulative, len(g.cum))
	for s, c := range g.cum {
		cc := *c
		if c.hist != nil {
			cc.hist = append([]uint64(nil), c.hist...)
		}
		out[s] = cc
	}
	return out
}