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

package metrics

import (
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
)

type Labels map[string]string

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

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

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) 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])
	}
	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 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
		}
		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)
}