lingniu-vehicle-ingest/go/vehicle-gateway/internal/metrics/metrics.go

package metrics

import (
	"fmt"
	"math"
	"net/http"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"
)

var processStartedAt = time.Now()

type Labels map[string]string

type Registry struct {
	mu         sync.RWMutex
	counters   map[string]map[string]float64
	gauges     map[string]map[string]float64
	histograms map[string]*histogramFamily
}

type histogramFamily struct {
	buckets map[string]map[float64]float64
	counts  map[string]float64
	sums    map[string]float64
}

func NewRegistry() *Registry {
	return &Registry{
		counters:   map[string]map[string]float64{},
		gauges:     map[string]map[string]float64{},
		histograms: map[string]*histogramFamily{},
	}
}

func RegisterServiceInfo(registry *Registry, service string) {
	if registry == nil {
		return
	}
	service = strings.TrimSpace(service)
	if service == "" {
		service = "vehicle-service"
	}
	registry.SetGauge("vehicle_service_info", Labels{"service": service}, 1)
	RecordProcessRuntime(registry)
}

func RecordProcessRuntime(registry *Registry) {
	if registry == nil {
		return
	}
	var mem runtime.MemStats
	runtime.ReadMemStats(&mem)
	registry.SetGauge("vehicle_process_start_time_unix_seconds", nil, float64(processStartedAt.Unix()))
	registry.SetGauge("vehicle_process_uptime_seconds", nil, time.Since(processStartedAt).Seconds())
	registry.SetGauge("vehicle_process_heap_alloc_bytes", nil, float64(mem.HeapAlloc))
	registry.SetGauge("vehicle_process_heap_sys_bytes", nil, float64(mem.HeapSys))
	registry.SetGauge("vehicle_process_goroutines", nil, float64(runtime.NumGoroutine()))
}

func RegisterKafkaConsumerInfo(registry *Registry, service string, group string, topics []string) {
	if registry == nil {
		return
	}
	service = strings.TrimSpace(service)
	if service == "" {
		service = "vehicle-service"
	}
	group = strings.TrimSpace(group)
	if group == "" {
		return
	}
	for _, topic := range topics {
		topic = strings.TrimSpace(topic)
		if topic == "" {
			continue
		}
		registry.SetGauge("vehicle_kafka_consumer_info", Labels{
			"service": service,
			"group":   group,
			"topic":   topic,
		}, 1)
	}
}

func (r *Registry) IncCounter(name string, labels Labels) {
	r.AddCounter(name, labels, 1)
}

func (r *Registry) AddCounter(name string, labels Labels, value float64) {
	name = strings.TrimSpace(name)
	if r == nil || name == "" || value == 0 {
		return
	}
	key := labelsKey(labels)
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.counters[name] == nil {
		r.counters[name] = map[string]float64{}
	}
	r.counters[name][key] += value
}

func (r *Registry) SetGauge(name string, labels Labels, value float64) {
	name = strings.TrimSpace(name)
	if r == nil || name == "" {
		return
	}
	key := labelsKey(labels)
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.gauges[name] == nil {
		r.gauges[name] = map[string]float64{}
	}
	r.gauges[name][key] = value
}

func (r *Registry) AddGauge(name string, labels Labels, value float64) {
	name = strings.TrimSpace(name)
	if r == nil || name == "" || value == 0 {
		return
	}
	key := labelsKey(labels)
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.gauges[name] == nil {
		r.gauges[name] = map[string]float64{}
	}
	r.gauges[name][key] += value
}

func (r *Registry) SetKafkaLag(name string, topic string, partition int, offset int64, highWaterMark int64) {
	lag := highWaterMark - offset - 1
	if lag < 0 {
		lag = 0
	}
	r.SetGauge(name, Labels{
		"topic":     topic,
		"partition": strconv.Itoa(partition),
	}, float64(lag))
}

func RecordLastActivity(registry *Registry, name string, labels Labels) {
	if registry == nil {
		return
	}
	registry.SetGauge(name, labels, float64(time.Now().Unix()))
}

func (r *Registry) ObserveHistogram(name string, labels Labels, buckets []float64, value float64) {
	name = strings.TrimSpace(name)
	if r == nil || name == "" {
		return
	}
	key := labelsKey(labels)
	r.mu.Lock()
	defer r.mu.Unlock()
	family := r.histograms[name]
	if family == nil {
		family = &histogramFamily{
			buckets: map[string]map[float64]float64{},
			counts:  map[string]float64{},
			sums:    map[string]float64{},
		}
		r.histograms[name] = family
	}
	if family.buckets[key] == nil {
		family.buckets[key] = map[float64]float64{}
	}
	normalizedBuckets := append([]float64(nil), buckets...)
	sort.Float64s(normalizedBuckets)
	for _, bucket := range normalizedBuckets {
		if value <= bucket {
			family.buckets[key][bucket]++
		} else if _, ok := family.buckets[key][bucket]; !ok {
			family.buckets[key][bucket] = 0
		}
	}
	family.buckets[key][math.Inf(1)]++
	family.counts[key]++
	family.sums[key] += value
}

func (r *Registry) Render() string {
	if r == nil {
		return ""
	}
	r.mu.RLock()
	defer r.mu.RUnlock()

	var names []string
	for name := range r.counters {
		names = append(names, name)
	}
	sort.Strings(names)

	var b strings.Builder
	for _, name := range names {
		renderMetricFamily(&b, name, "counter", r.counters[name])
	}

	names = names[:0]
	for name := range r.gauges {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		renderMetricFamily(&b, name, "gauge", r.gauges[name])
	}

	names = names[:0]
	for name := range r.histograms {
		names = append(names, name)
	}
	sort.Strings(names)
	for _, name := range names {
		renderHistogramFamily(&b, name, r.histograms[name])
	}
	return b.String()
}

func renderMetricFamily(b *strings.Builder, name string, metricType string, values map[string]float64) {
	b.WriteString("# TYPE ")
	b.WriteString(name)
	b.WriteString(" ")
	b.WriteString(metricType)
	b.WriteString("\n")
	var keys []string
	for key := range values {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	for _, key := range keys {
		b.WriteString(name)
		if key != "" {
			b.WriteString("{")
			b.WriteString(key)
			b.WriteString("}")
		}
		b.WriteString(" ")
		b.WriteString(formatNumber(values[key]))
		b.WriteString("\n")
	}
}

func renderHistogramFamily(b *strings.Builder, name string, family *histogramFamily) {
	if family == nil {
		return
	}
	b.WriteString("# TYPE ")
	b.WriteString(name)
	b.WriteString(" histogram\n")
	var keys []string
	for key := range family.counts {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	for _, key := range keys {
		var buckets []float64
		for bucket := range family.buckets[key] {
			buckets = append(buckets, bucket)
		}
		sort.Float64s(buckets)
		for _, bucket := range buckets {
			renderHistogramMetric(b, name+"_bucket", key, "le", bucketLabel(bucket), family.buckets[key][bucket])
		}
		renderMetricSample(b, name+"_count", key, family.counts[key])
		renderMetricSample(b, name+"_sum", key, family.sums[key])
	}
}

func renderHistogramMetric(b *strings.Builder, name string, key string, extraLabel string, extraValue string, value float64) {
	b.WriteString(name)
	labels := mergeLabelKey(key, extraLabel, extraValue)
	if labels != "" {
		b.WriteString("{")
		b.WriteString(labels)
		b.WriteString("}")
	}
	b.WriteString(" ")
	b.WriteString(formatNumber(value))
	b.WriteString("\n")
}

func renderMetricSample(b *strings.Builder, name string, key string, value float64) {
	b.WriteString(name)
	if key != "" {
		b.WriteString("{")
		b.WriteString(key)
		b.WriteString("}")
	}
	b.WriteString(" ")
	b.WriteString(formatNumber(value))
	b.WriteString("\n")
}

func mergeLabelKey(key string, extraLabel string, extraValue string) string {
	extra := fmt.Sprintf(`%s="%s"`, sanitizeLabelName(extraLabel), escapeLabelValue(extraValue))
	if key == "" {
		return extra
	}
	return extra + "," + key
}

func bucketLabel(value float64) string {
	if math.IsInf(value, 1) {
		return "+Inf"
	}
	return formatNumber(value)
}

func NewHandler(registry *Registry) http.Handler {
	if registry == nil {
		registry = NewRegistry()
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodGet {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		if strings.Trim(r.URL.Path, "/") != "metrics" {
			http.NotFound(w, r)
			return
		}
		RecordProcessRuntime(registry)
		w.Header().Set("Content-Type", "text/plain; version=0.0.4; charset=utf-8")
		_, _ = w.Write([]byte(registry.Render()))
	})
}

func labelsKey(labels Labels) string {
	if len(labels) == 0 {
		return ""
	}
	keys := make([]string, 0, len(labels))
	for key := range labels {
		keys = append(keys, key)
	}
	sort.Strings(keys)
	parts := make([]string, 0, len(keys))
	for _, key := range keys {
		value := labels[key]
		parts = append(parts, fmt.Sprintf(`%s="%s"`, sanitizeLabelName(key), escapeLabelValue(value)))
	}
	return strings.Join(parts, ",")
}

func sanitizeLabelName(value string) string {
	value = strings.TrimSpace(value)
	if value == "" {
		return "label"
	}
	var b strings.Builder
	for i, r := range value {
		ok := r == '_' || (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (i > 0 && r >= '0' && r <= '9')
		if ok {
			b.WriteRune(r)
		} else {
			b.WriteByte('_')
		}
	}
	return b.String()
}

func escapeLabelValue(value string) string {
	value = strings.ReplaceAll(value, `\`, `\\`)
	value = strings.ReplaceAll(value, "\n", `\n`)
	return strings.ReplaceAll(value, `"`, `\"`)
}

func formatNumber(value float64) string {
	if value == float64(int64(value)) {
		return fmt.Sprintf("%d", int64(value))
	}
	return fmt.Sprintf("%g", value)
}