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 @@ funcRet = fmt.Sprintf(" returning %T", v) } } +// 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 @@ -193,6 +198,18 @@ func() *expvar.Map { })(), "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",