From 7b4960316bed5fea4ee55a6d42443daf8f8c1125 Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Sat, 19 Mar 2022 23:16:06 -0700
Subject: [PATCH] tsweb: add PrometheusVar, for vars that want to output varz
 themselves.

This enables the infrequent use of more complex Prometheus types, such as
timeseries with high/irregular label cardinality, without needing to
discover and implement generic abstracted type like LabelMap for each one.

Signed-off-by: David Anderson <danderson@tailscale.com>
---
 tsweb/tsweb.go      | 12 ++++++++++++
 tsweb/tsweb_test.go | 18 ++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index 59ebbf1e2..63b2e6071 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -349,6 +349,15 @@ func Error(code int, msg string, err error) HTTPError {
 	return HTTPError{Code: code, Msg: msg, Err: err}
 }
 
+// PrometheusVar is a value that knows how to format itself into
+// Prometheus metric syntax.
+type PrometheusVar interface {
+	// WritePrometheus writes the value of the var to w, in Prometheus
+	// metric syntax. All variables names written out must start with
+	// prefix (or write out a single variable named exactly prefix)
+	WritePrometheus(w io.Writer, prefix string)
+}
+
 // WritePrometheusExpvar writes kv to w in Prometheus metrics format.
 //
 // See VarzHandler for conventions. This is exported primarily for
@@ -379,6 +388,9 @@ func writePromExpVar(w io.Writer, prefix string, kv expvar.KeyValue) {
 	name := prefix + key
 
 	switch v := kv.Value.(type) {
+	case PrometheusVar:
+		v.WritePrometheus(w, name)
+		return
 	case *expvar.Int:
 		if typ == "" {
 			typ = "counter"
diff --git a/tsweb/tsweb_test.go b/tsweb/tsweb_test.go
index c7469f5da..b5c081b09 100644
--- a/tsweb/tsweb_test.go
+++ b/tsweb/tsweb_test.go
@@ -9,6 +9,8 @@
 	"context"
 	"errors"
 	"expvar"
+	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"net/http/httptest"
@@ -469,6 +471,12 @@ func() *expvar.Map {
 			expvar.Func(func() any { return 123 }),
 			"num_goroutines 123\n",
 		},
+		{
+			"var_that_exports_itself",
+			"custom_var",
+			promWriter{},
+			"custom_var_value 42\n",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -534,3 +542,13 @@ func (expvarAdapter) String() string { return "{}" } // expvar JSON; unused in t
 func (a expvarAdapter) PrometheusMetricsReflectRoot() any {
 	return a.st
 }
+
+type promWriter struct{}
+
+func (promWriter) WritePrometheus(w io.Writer, prefix string) {
+	fmt.Fprintf(w, "%s_value 42\n", prefix)
+}
+
+func (promWriter) String() string {
+	return ""
+}