146 lines
3.9 KiB
Go
146 lines
3.9 KiB
Go
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()
|
|
}
|