orca/internal/transport/metrics.go

package transport

import (
	"fmt"
	"io"
	"strconv"
	"sync"
	"sync/atomic"
)

type metricDef struct {
	help string
	typ  string
}

var metricOrder = []string{
	"txns_applied_total",
	"txns_drifted_total",
	"drifts_remediated_total",
	"peers_total",
	"nodes_total",
	"allocs_total",
}

var metricMeta = map[string]metricDef{
	"txns_applied_total":      {"Total transactions applied", "counter"},
	"txns_drifted_total":      {"Total transactions drifted", "counter"},
	"drifts_remediated_total": {"Total drifts remediated", "counter"},
	"peers_total":             {"Current peer count", "gauge"},
	"nodes_total":             {"Current node count", "gauge"},
	"allocs_total":            {"Current allocation count", "gauge"},
}

type Metrics struct {
	cmu      sync.RWMutex
	counters map[string]*atomic.Int64

	gmu    sync.RWMutex
	gauges map[string]float64
}

func NewMetrics() *Metrics {
	m := &Metrics{
		counters: make(map[string]*atomic.Int64),
		gauges:   make(map[string]float64),
	}
	for name, meta := range metricMeta {
		if meta.typ == "counter" {
			m.counters[name] = new(atomic.Int64)
		} else {
			m.gauges[name] = 0
		}
	}
	return m
}

func (m *Metrics) counter(name string) *atomic.Int64 {
	m.cmu.RLock()
	c := m.counters[name]
	m.cmu.RUnlock()
	if c != nil {
		return c
	}
	m.cmu.Lock()
	defer m.cmu.Unlock()
	if c = m.counters[name]; c != nil {
		return c
	}
	c = new(atomic.Int64)
	m.counters[name] = c
	return c
}

func (m *Metrics) IncCounter(name string) {
	m.counter(name).Add(1)
}

func (m *Metrics) AddCounter(name string, delta int64) {
	m.counter(name).Add(delta)
}

func (m *Metrics) SetGauge(name string, value float64) {
	m.gmu.Lock()
	m.gauges[name] = value
	m.gmu.Unlock()
}

func (m *Metrics) WritePrometheus(w io.Writer) error {
	for _, name := range metricOrder {
		meta, ok := metricMeta[name]
		if !ok {
			continue
		}
		if _, err := fmt.Fprintf(w, "# HELP %s %s\n", name, meta.help); err != nil {
			return err
		}
		if _, err := fmt.Fprintf(w, "# TYPE %s %s\n", name, meta.typ); err != nil {
			return err
		}
		switch meta.typ {
		case "counter":
			m.cmu.RLock()
			c := m.counters[name]
			m.cmu.RUnlock()
			var v int64
			if c != nil {
				v = c.Load()
			}
			if _, err := fmt.Fprintf(w, "%s %d\n", name, v); err != nil {
				return err
			}
		case "gauge":
			m.gmu.RLock()
			v := m.gauges[name]
			m.gmu.RUnlock()
			if _, err := fmt.Fprintf(w, "%s %s\n", name, strconv.FormatFloat(v, 'g', -1, 64)); err != nil {
				return err
			}
		}
	}
	return nil
}