reforger_crawler_main/metrics/prometheus.go

package metrics

import (
	"bufio"
	"io"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"time"

	"gitea.tbdevent.eu/TBD/reforger_crawler_main/initializers"
)

// WritePrometheus writes all cumulative series in the Prometheus text format
// (version 0.0.4). External metrics are exposed as <service>_<name> with an
// "instance" label; counters are totals since this process started.
func WritePrometheus(w io.Writer) {
	bw := bufio.NewWriter(w)
	defer bw.Flush()

	cum := defaultAgg.snapshotCum()
	type entry struct {
		s *Series
		c cumulative
	}
	byName := map[string][]entry{}
	meta := map[string]*Series{}
	for s, c := range cum {
		name := s.Name
		if s.Service != MainService {
			name = s.Service + "_" + s.Name
		}
		byName[name] = append(byName[name], entry{s, c})
		if meta[name] == nil {
			meta[name] = s
		}
	}
	names := make([]string, 0, len(byName))
	for n := range byName {
		names = append(names, n)
	}
	sort.Strings(names)

	for _, name := range names {
		s := meta[name]
		writeHeader(bw, name, s.Help, promType(s.Kind))
		entries := byName[name]
		sort.Slice(entries, func(i, j int) bool {
			if entries[i].s.Instance != entries[j].s.Instance {
				return entries[i].s.Instance < entries[j].s.Instance
			}
			return entries[i].s.LabelsJSON < entries[j].s.LabelsJSON
		})
		for _, e := range entries {
			labels := e.s.Labels
			if e.s.Service != MainService {
				labels = withLabel(labels, "instance", e.s.Instance)
			}
			switch e.s.Kind {
			case Counter:
				writeSample(bw, name, labels, e.c.sum)
			case Gauge:
				writeSample(bw, name, labels, e.c.last)
			case Summary:
				writeSample(bw, name+"_sum", labels, e.c.sum)
				writeSample(bw, name+"_count", labels, float64(e.c.count))
			case Histogram:
				var acc uint64
				for i, b := range e.s.Bounds {
					if i < len(e.c.hist) {
						acc += e.c.hist[i]
					}
					writeSample(bw, name+"_bucket", withLabel(labels, "le", formatFloat(b)), float64(acc))
				}
				writeSample(bw, name+"_bucket", withLabel(labels, "le", "+Inf"), float64(e.c.count))
				writeSample(bw, name+"_sum", labels, e.c.sum)
				writeSample(bw, name+"_count", labels, float64(e.c.count))
			}
		}
	}

	up := ServiceUp()
	if len(up) > 0 {
		writeHeader(bw, "reforger_service_up", "Whether a reporting service instance is sending heartbeats", "gauge")
		keys := make([][2]string, 0, len(up))
		for k := range up {
			keys = append(keys, k)
		}
		sort.Slice(keys, func(i, j int) bool { return keys[i][0]+keys[i][1] < keys[j][0]+keys[j][1] })
		for _, k := range keys {
			v := 0.0
			if up[k] {
				v = 1
			}
			writeSample(bw, "reforger_service_up", map[string]string{"service": k[0], "instance": k[1]}, v)
		}
	}

	var ms runtime.MemStats
	runtime.ReadMemStats(&ms)
	writeHeader(bw, "go_goroutines", "Number of goroutines", "gauge")
	writeSample(bw, "go_goroutines", nil, float64(runtime.NumGoroutine()))
	writeHeader(bw, "go_memstats_heap_alloc_bytes", "Heap bytes allocated and in use", "gauge")
	writeSample(bw, "go_memstats_heap_alloc_bytes", nil, float64(ms.HeapAlloc))
	writeHeader(bw, "go_memstats_sys_bytes", "Bytes obtained from the OS", "gauge")
	writeSample(bw, "go_memstats_sys_bytes", nil, float64(ms.Sys))
	writeHeader(bw, "go_gc_cycles_total", "Completed GC cycles", "counter")
	writeSample(bw, "go_gc_cycles_total", nil, float64(ms.NumGC))
	writeHeader(bw, "process_start_time_seconds", "Process start time", "gauge")
	writeSample(bw, "process_start_time_seconds", nil, float64(startedAt.Unix()))
	writeHeader(bw, "process_uptime_seconds", "Process uptime", "gauge")
	writeSample(bw, "process_uptime_seconds", nil, time.Since(startedAt).Seconds())

	if sqlDB, err := initializers.DB.DB(); err == nil {
		st := sqlDB.Stats()
		writeHeader(bw, "db_open_connections", "Open database connections", "gauge")
		writeSample(bw, "db_open_connections", nil, float64(st.OpenConnections))
		writeHeader(bw, "db_in_use_connections", "Database connections in use", "gauge")
		writeSample(bw, "db_in_use_connections", nil, float64(st.InUse))
		writeHeader(bw, "db_wait_count_total", "Waits for a database connection", "counter")
		writeSample(bw, "db_wait_count_total", nil, float64(st.WaitCount))
		writeHeader(bw, "db_wait_seconds_total", "Time spent waiting for a database connection", "counter")
		writeSample(bw, "db_wait_seconds_total", nil, st.WaitDuration.Seconds())
	}
}

func promType(k Kind) string {
	switch k {
	case Counter:
		return "counter"
	case Histogram:
		return "histogram"
	case Summary:
		return "summary"
	}
	return "gauge"
}

func withLabel(labels map[string]string, k, v string) map[string]string {
	out := make(map[string]string, len(labels)+1)
	for lk, lv := range labels {
		out[lk] = lv
	}
	out[k] = v
	return out
}

func formatFloat(v float64) string { return strconv.FormatFloat(v, 'g', -1, 64) }

var helpEscaper = strings.NewReplacer(`\`, `\\`, "\n", `\n`)
var valueEscaper = strings.NewReplacer(`\`, `\\`, "\n", `\n`, `"`, `\"`)

func writeHeader(w *bufio.Writer, name, help, typ string) {
	if help != "" {
		w.WriteString("# HELP " + name + " " + helpEscaper.Replace(help) + "\n")
	}
	w.WriteString("# TYPE " + name + " " + typ + "\n")
}

func writeSample(w *bufio.Writer, name string, labels map[string]string, v float64) {
	w.WriteString(name)
	if len(labels) > 0 {
		keys := make([]string, 0, len(labels))
		for k := range labels {
			keys = append(keys, k)
		}
		sort.Strings(keys)
		w.WriteByte('{')
		for i, k := range keys {
			if i > 0 {
				w.WriteByte(',')
			}
			w.WriteString(k + `="` + valueEscaper.Replace(labels[k]) + `"`)
		}
		w.WriteByte('}')
	}
	w.WriteByte(' ')
	w.WriteString(formatFloat(v))
	w.WriteByte('\n')
}