metrics, tsweb/varz: add multi-label map metrics

Updates tailscale/corp#18640

Change-Id: Ia9ae25956038e9d3266ea165537ac6f02485b74c
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

4 files changed, 408 insertions, 0 deletions

diff --git a/metrics/multilabelmap.go b/metrics/multilabelmap.go
new file mode 100644
index 000000000..9edf1aec0
--- /dev/null
+++ b/metrics/multilabelmap.go
@@ -0,0 +1,259 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package metrics
+
+import (
+	"expvar"
+	"fmt"
+	"io"
+	"reflect"
+	"sort"
+	"strings"
+	"sync"
+)
+
+// MultiLabelMap is a struct-value-to-Var map variable that satisfies the
+// [expvar.Var] interface but also allows for multiple Prometheus labels to be
+// associated with each value.
+//
+// T must be a struct type with only string fields. The struct field names
+// (lowercased) are used as the labels, unless a "prom" struct tag is present.
+// The struct fields must all be strings, and the string values must be valid
+// Prometheus label values without requiring quoting.
+type MultiLabelMap[T comparable] struct {
+	Type string // optional Prometheus type ("counter", "gauge")
+	Help string // optional Prometheus help string
+
+	m sync.Map // map[T]expvar.Var
+
+	mu     sync.RWMutex
+	sorted []labelsAndValue[T] // by labels string, to match expvar.Map + for aesthetics in output
+}
+
+// NewMultiLabelMap creates and publishes (via expvar.Publish) a new
+// MultiLabelMap[T] variable with the given name and returns it.
+func NewMultiLabelMap[T comparable](name string, promType, helpText string) *MultiLabelMap[T] {
+	m := &MultiLabelMap[T]{
+		Type: promType,
+		Help: helpText,
+	}
+	expvar.Publish(name, m)
+	return m
+}
+
+type labelsAndValue[T comparable] struct {
+	key    T
+	labels string // Prometheus-formatted {label="value",label="value"} string
+	val    expvar.Var
+}
+
+// labelString returns a Prometheus-formatted label string for the given key.
+func labelString(k any) string {
+	rv := reflect.ValueOf(k)
+	t := rv.Type()
+	if t.Kind() != reflect.Struct {
+		panic(fmt.Sprintf("MultiLabelMap must use keys of type struct; got %v", t))
+	}
+
+	var sb strings.Builder
+	sb.WriteString("{")
+
+	for i := 0; i < t.NumField(); i++ {
+		if i > 0 {
+			sb.WriteString(",")
+		}
+		ft := t.Field(i)
+		label := ft.Tag.Get("prom")
+		if label == "" {
+			label = strings.ToLower(ft.Name)
+		}
+		fv := rv.Field(i)
+		switch fv.Kind() {
+		case reflect.String:
+			fmt.Fprintf(&sb, "%s=%q", label, fv.String())
+		case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
+			fmt.Fprintf(&sb, "%s=\"%d\"", label, fv.Int())
+		case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
+			fmt.Fprintf(&sb, "%s=\"%d\"", label, fv.Uint())
+		case reflect.Bool:
+			fmt.Fprintf(&sb, "%s=\"%v\"", label, fv.Bool())
+		default:
+			panic(fmt.Sprintf("MultiLabelMap key field %q has unsupported type %v", ft.Name, fv.Type()))
+		}
+	}
+	sb.WriteString("}")
+	return sb.String()
+}
+
+// KeyValue represents a single entry in a [MultiLabelMap].
+type KeyValue[T comparable] struct {
+	Key   T
+	Value expvar.Var
+}
+
+func (v *MultiLabelMap[T]) String() string {
+	return `"MultiLabelMap"`
+}
+
+// WritePrometheus writes v to w in Prometheus exposition format.
+// The name argument is the metric name.
+func (v *MultiLabelMap[T]) WritePrometheus(w io.Writer, name string) {
+	if v.Type != "" {
+		io.WriteString(w, "# TYPE ")
+		io.WriteString(w, name)
+		io.WriteString(w, " ")
+		io.WriteString(w, v.Type)
+		io.WriteString(w, "\n")
+	}
+	if v.Help != "" {
+		io.WriteString(w, "# HELP ")
+		io.WriteString(w, name)
+		io.WriteString(w, " ")
+		io.WriteString(w, v.Help)
+		io.WriteString(w, "\n")
+	}
+	v.mu.RLock()
+	defer v.mu.RUnlock()
+
+	for _, kv := range v.sorted {
+		io.WriteString(w, name)
+		io.WriteString(w, kv.labels)
+		switch v := kv.val.(type) {
+		case *expvar.Int:
+			fmt.Fprintf(w, " %d\n", v.Value())
+		case *expvar.Float:
+			fmt.Fprintf(w, " %v\n", v.Value())
+		default:
+			fmt.Fprintf(w, " %s\n", kv.val)
+		}
+	}
+}
+
+// Init removes all keys from the map.
+//
+// Think of it as "Reset", but it's named Init to match expvar.Map.Init.
+func (v *MultiLabelMap[T]) Init() *MultiLabelMap[T] {
+	v.mu.Lock()
+	defer v.mu.Unlock()
+	v.sorted = nil
+	v.m.Range(func(k, _ any) bool {
+		v.m.Delete(k)
+		return true
+	})
+	return v
+}
+
+// addKeyLocked updates the sorted list of keys in v.keys.
+//
+// v.mu must be held.
+func (v *MultiLabelMap[T]) addKeyLocked(key T, val expvar.Var) {
+	ls := labelString(key)
+
+	ent := labelsAndValue[T]{key, ls, val}
+	// Using insertion sort to place key into the already-sorted v.keys.
+	i := sort.Search(len(v.sorted), func(i int) bool {
+		return v.sorted[i].labels >= ls
+	})
+	if i >= len(v.sorted) {
+		v.sorted = append(v.sorted, ent)
+	} else if v.sorted[i].key == key {
+		v.sorted[i].val = val
+	} else {
+		var zero labelsAndValue[T]
+		v.sorted = append(v.sorted, zero)
+		copy(v.sorted[i+1:], v.sorted[i:])
+		v.sorted[i] = ent
+	}
+}
+
+// Get returns the expvar for the given key, or nil if it doesn't exist.
+func (v *MultiLabelMap[T]) Get(key T) expvar.Var {
+	i, _ := v.m.Load(key)
+	av, _ := i.(expvar.Var)
+	return av
+}
+
+func newInt() expvar.Var   { return new(expvar.Int) }
+func newFloat() expvar.Var { return new(expvar.Float) }
+
+// getOrFill returns the expvar.Var for the given key, atomically creating it
+// once (for all callers) with fill if it doesn't exist.
+func (v *MultiLabelMap[T]) getOrFill(key T, fill func() expvar.Var) expvar.Var {
+	if v := v.Get(key); v != nil {
+		return v
+	}
+
+	v.mu.Lock()
+	defer v.mu.Unlock()
+
+	if v := v.Get(key); v != nil {
+		return v
+	}
+	nv := fill()
+	v.addKeyLocked(key, nv)
+	v.m.Store(key, nv)
+	return nv
+}
+
+// Set sets key to val.
+//
+// This is not optimized for highly concurrent usage; it's presumed to only be
+// used rarely, at startup.
+func (v *MultiLabelMap[T]) Set(key T, val expvar.Var) {
+	v.mu.Lock()
+	defer v.mu.Unlock()
+	v.addKeyLocked(key, val)
+	v.m.Store(key, val)
+}
+
+// Add adds delta to the *[expvar.Int] value stored under the given map key,
+// creating it if it doesn't exist yet.
+// It does nothing if key exists but is of the wrong type.
+func (v *MultiLabelMap[T]) Add(key T, delta int64) {
+	// Add to Int; ignore otherwise.
+	if iv, ok := v.getOrFill(key, newInt).(*expvar.Int); ok {
+		iv.Add(delta)
+	}
+}
+
+// Add adds delta to the *[expvar.Float] value stored under the given map key,
+// creating it if it doesn't exist yet.
+// It does nothing if key exists but is of the wrong type.
+func (v *MultiLabelMap[T]) AddFloat(key T, delta float64) {
+	// Add to Float; ignore otherwise.
+	if iv, ok := v.getOrFill(key, newFloat).(*expvar.Float); ok {
+		iv.Add(delta)
+	}
+}
+
+// Delete deletes the given key from the map.
+//
+// This is not optimized for highly concurrent usage; it's presumed to only be
+// used rarely, at startup.
+func (v *MultiLabelMap[T]) Delete(key T) {
+	ls := labelString(key)
+
+	v.mu.Lock()
+	defer v.mu.Unlock()
+
+	// Using insertion sort to place key into the already-sorted v.keys.
+	i := sort.Search(len(v.sorted), func(i int) bool {
+		return v.sorted[i].labels >= ls
+	})
+	if i < len(v.sorted) && v.sorted[i].key == key {
+		v.sorted = append(v.sorted[:i], v.sorted[i+1:]...)
+		v.m.Delete(key)
+	}
+}
+
+// Do calls f for each entry in the map.
+// The map is locked during the iteration,
+// but existing entries may be concurrently updated.
+func (v *MultiLabelMap[T]) Do(f func(KeyValue[T])) {
+	v.mu.RLock()
+	defer v.mu.RUnlock()
+	for _, e := range v.sorted {
+		f(KeyValue[T]{e.key, e.val})
+	}
+}
diff --git a/metrics/multilabelmap_test.go b/metrics/multilabelmap_test.go
new file mode 100644
index 000000000..9a1340a3c
--- /dev/null
+++ b/metrics/multilabelmap_test.go
@@ -0,0 +1,121 @@
+// Copyright (c) Tailscale Inc & AUTHORS
+// SPDX-License-Identifier: BSD-3-Clause
+
+package metrics
+
+import (
+	"bytes"
+	"fmt"
+	"io"
+	"testing"
+)
+
+type L2 struct {
+	Foo string `prom:"foo"`
+	Bar string `prom:"bar"`
+}
+
+func TestMultilabelMap(t *testing.T) {
+	m := new(MultiLabelMap[L2])
+	m.Add(L2{"a", "b"}, 2)
+	m.Add(L2{"b", "c"}, 4)
+	m.Add(L2{"b", "b"}, 3)
+	m.Add(L2{"a", "a"}, 1)
+
+	cur := func() string {
+		var buf bytes.Buffer
+		m.Do(func(kv KeyValue[L2]) {
+			if buf.Len() > 0 {
+				buf.WriteString(",")
+			}
+			fmt.Fprintf(&buf, "%s/%s=%v", kv.Key.Foo, kv.Key.Bar, kv.Value)
+		})
+		return buf.String()
+	}
+
+	if g, w := cur(), "a/a=1,a/b=2,b/b=3,b/c=4"; g != w {
+		t.Errorf("got %q; want %q", g, w)
+	}
+
+	var buf bytes.Buffer
+	m.WritePrometheus(&buf, "metricname")
+	const want = `metricname{foo="a",bar="a"} 1
+metricname{foo="a",bar="b"} 2
+metricname{foo="b",bar="b"} 3
+metricname{foo="b",bar="c"} 4
+`
+	if got := buf.String(); got != want {
+		t.Errorf("promtheus output = %q; want %q", got, want)
+	}
+
+	m.Delete(L2{"b", "b"})
+
+	if g, w := cur(), "a/a=1,a/b=2,b/c=4"; g != w {
+		t.Errorf("got %q; want %q", g, w)
+	}
+
+	allocs := testing.AllocsPerRun(1000, func() {
+		m.Add(L2{"a", "a"}, 1)
+	})
+	if allocs > 0 {
+		t.Errorf("allocs = %v; want 0", allocs)
+	}
+	m.Init()
+	if g, w := cur(), ""; g != w {
+		t.Errorf("got %q; want %q", g, w)
+	}
+
+	writeAllocs := testing.AllocsPerRun(1000, func() {
+		m.WritePrometheus(io.Discard, "test")
+	})
+	if writeAllocs > 0 {
+		t.Errorf("writeAllocs = %v; want 0", writeAllocs)
+	}
+}
+
+func TestMultiLabelMapTypes(t *testing.T) {
+	type LabelTypes struct {
+		S string
+		B bool
+		I int
+		U uint
+	}
+
+	m := new(MultiLabelMap[LabelTypes])
+	m.Type = "counter"
+	m.Help = "some good stuff"
+	m.Add(LabelTypes{"a", true, -1, 2}, 3)
+	var buf bytes.Buffer
+	m.WritePrometheus(&buf, "metricname")
+	const want = `# TYPE metricname counter
+# HELP metricname some good stuff
+metricname{s="a",b="true",i="-1",u="2"} 3
+`
+	if got := buf.String(); got != want {
+		t.Errorf("got %q; want %q", got, want)
+	}
+
+	writeAllocs := testing.AllocsPerRun(1000, func() {
+		m.WritePrometheus(io.Discard, "test")
+	})
+	if writeAllocs > 0 {
+		t.Errorf("writeAllocs = %v; want 0", writeAllocs)
+	}
+}
+
+func BenchmarkMultiLabelWriteAllocs(b *testing.B) {
+	b.ReportAllocs()
+
+	m := new(MultiLabelMap[L2])
+	m.Add(L2{"a", "b"}, 2)
+	m.Add(L2{"b", "c"}, 4)
+	m.Add(L2{"b", "b"}, 3)
+	m.Add(L2{"a", "a"}, 1)
+
+	var w io.Writer = io.Discard
+
+	b.ResetTimer()
+	for range b.N {
+		m.WritePrometheus(w, "test")
+	}
+}
diff --git a/tsweb/varz/varz.go b/tsweb/varz/varz.go
index 243cc0782..e3e52e031 100644
--- a/tsweb/varz/varz.go
+++ b/tsweb/varz/varz.go
@@ -133,6 +133,9 @@ func writePromExpVar(w io.Writer, prefix string, kv expvar.KeyValue) {
 			writePromExpVar(w, name+"_", kv)
 		})
 		return
+	case PrometheusWriter:
+		v.WritePrometheus(w, name)
+		return
 	case PrometheusMetricsReflectRooter:
 		root := v.PrometheusMetricsReflectRoot()
 		rv := reflect.ValueOf(root)
@@ -233,6 +236,14 @@ func writePromExpVar(w io.Writer, prefix string, kv expvar.KeyValue) {
 	}
 }
 
+// PrometheusWriter is the interface implemented by metrics that can write
+// themselves into Prometheus exposition format.
+//
+// As of 2024-03-25, this is only *metrics.MultiLabelMap.
+type PrometheusWriter interface {
+	WritePrometheus(w io.Writer, name string)
+}
+
 var sortedKVsPool = &sync.Pool{New: func() any { return new(sortedKVs) }}
 
 // sortedKV is a KeyValue with a sort key.
diff --git a/tsweb/varz/varz_test.go b/tsweb/varz/varz_test.go
index 6105f571b..7e094b0e7 100644
--- a/tsweb/varz/varz_test.go
+++ b/tsweb/varz/varz_test.go
@@ -25,6 +25,11 @@ func TestVarzHandler(t *testing.T) {
 	half := new(expvar.Float)
 	half.Set(0.5)
 
+	type L2 struct {
+		Foo string `prom:"foo"`
+		Bar string `prom:"bar"`
+	}
+
 	tests := []struct {
 		name string
 		k    string // key name
@@ -194,6 +199,18 @@ func TestVarzHandler(t *testing.T) {
 			"foo{label=\"a\"} 1\n",
 		},
 		{
+			"metrics_multilabel_map",
+			"foo",
+			(func() *metrics.MultiLabelMap[L2] {
+				m := new(metrics.MultiLabelMap[L2])
+				m.Add(L2{"a", "b"}, 1)
+				m.Add(L2{"c", "d"}, 2)
+				return m
+			})(),
+			"foo{foo=\"a\",bar=\"b\"} 1\n" +
+				"foo{foo=\"c\",bar=\"d\"} 2\n",
+		},
+		{
 			"expvar_label_map",
 			"counter_labelmap_keyname_m",
 			func() *expvar.Map {