step one

internal/metrics/prom.go

@@ -0,0 +1,146 @@
+package metrics
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+)
+
+// Prom is a minimal Sink that exposes counters and gauges in Prometheus's
+// text exposition format. No external dependencies.
+//
+// The Prom sink supports labelled counters, sums (for durations and byte
+// counts), and labelled gauges. Histograms are intentionally omitted; if
+// they are needed later, swap Prom for an OTel-based sink.
+type Prom struct {
+	mu       sync.Mutex
+	counters map[string]map[string]float64 // name → label-key → value
+	gauges   map[string]map[string]float64 // name → label-key → value
+}
+
+// NewProm returns an empty Prom sink.
+func NewProm() *Prom {
+	return &Prom{
+		counters: make(map[string]map[string]float64),
+		gauges:   make(map[string]map[string]float64),
+	}
+}
+
+// Prediction implements Sink.
+func (p *Prom) Prediction(profile string, d time.Duration, err error) {
+	status := "ok"
+	if err != nil {
+		status = "error"
+	}
+	labels := map[string]string{"profile": profile, "status": status}
+	p.incCounter("predictor_predictions_total", labels, 1)
+	p.incCounter("predictor_prediction_duration_seconds_sum", labels, d.Seconds())
+	p.incCounter("predictor_prediction_duration_seconds_count", labels, 1)
+}
+
+// Download implements Sink.
+func (p *Prom) Download(source string, d time.Duration, status string, bytes int64) {
+	labels := map[string]string{"source": source, "status": status}
+	p.incCounter("predictor_downloads_total", labels, 1)
+	p.incCounter("predictor_download_duration_seconds_sum", labels, d.Seconds())
+	p.incCounter("predictor_download_bytes_total", map[string]string{"source": source}, float64(bytes))
+}
+
+// ActiveEpoch implements Sink.
+func (p *Prom) ActiveEpoch(t time.Time) {
+	var v float64
+	if !t.IsZero() {
+		v = float64(t.Unix())
+	}
+	p.setGauge("predictor_active_dataset_epoch_seconds", map[string]string{}, v)
+}
+
+// ServeHTTP writes the metrics in Prometheus text exposition format.
+func (p *Prom) ServeHTTP(w http.ResponseWriter, _ *http.Request) {
+	w.Header().Set("Content-Type", "text/plain; version=0.0.4")
+	p.Write(w)
+}
+
+// Write writes the metrics in Prometheus exposition format to w.
+func (p *Prom) Write(w io.Writer) {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+
+	names := make([]string, 0, len(p.counters)+len(p.gauges))
+	for n := range p.counters {
+		names = append(names, n)
+	}
+	for n := range p.gauges {
+		names = append(names, n)
+	}
+	sort.Strings(names)
+
+	for _, name := range names {
+		if labels, ok := p.counters[name]; ok {
+			fmt.Fprintf(w, "# TYPE %s counter\n", name)
+			writeMetricFamily(w, name, labels)
+		}
+		if labels, ok := p.gauges[name]; ok {
+			fmt.Fprintf(w, "# TYPE %s gauge\n", name)
+			writeMetricFamily(w, name, labels)
+		}
+	}
+}
+
+func writeMetricFamily(w io.Writer, name string, labels map[string]float64) {
+	keys := make([]string, 0, len(labels))
+	for k := range labels {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	for _, key := range keys {
+		fmt.Fprintf(w, "%s%s %g\n", name, key, labels[key])
+	}
+}
+
+func (p *Prom) incCounter(name string, labels map[string]string, n float64) {
+	key := labelKey(labels)
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.counters[name] == nil {
+		p.counters[name] = make(map[string]float64)
+	}
+	p.counters[name][key] += n
+}
+
+func (p *Prom) setGauge(name string, labels map[string]string, v float64) {
+	key := labelKey(labels)
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if p.gauges[name] == nil {
+		p.gauges[name] = make(map[string]float64)
+	}
+	p.gauges[name][key] = v
+}
+
+// labelKey renders the labels into a Prometheus-format "{k1="v1",k2="v2"}"
+// suffix, empty if no labels.
+func labelKey(labels map[string]string) string {
+	if len(labels) == 0 {
+		return ""
+	}
+	keys := make([]string, 0, len(labels))
+	for k := range labels {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	var sb strings.Builder
+	sb.WriteByte('{')
+	for i, k := range keys {
+		if i > 0 {
+			sb.WriteByte(',')
+		}
+		fmt.Fprintf(&sb, "%s=%q", k, labels[k])
+	}
+	sb.WriteByte('}')
+	return sb.String()
+}