metrics, tsweb: add Distribution typeandrew/metrics-distribution

Signed-off-by: Andrew Dunham <andrew@du.nham.ca>
Change-Id: Ia3a87cccc35623dccdf74c0ff49f2785ced3c0df

4 files changed, 195 insertions, 1 deletions

diff --git a/metrics/metrics.go b/metrics/metrics.go
index bc19a1849..33d64b13e 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -6,7 +6,10 @@
 // Tailscale for monitoring.
 package metrics
 
-import "expvar"
+import (
+	"expvar"
+	"fmt"
+)
 
 // Set is a string-to-Var map variable that satisfies the expvar.Var
 // interface.
@@ -54,3 +57,70 @@ func (m *LabelMap) GetFloat(key string) *expvar.Float {
 func CurrentFDs() int {
 	return currentFDs()
 }
+
+// Distribution represents a set of values separated into individual "bins".
+//
+// Semantically, this is mapped by tsweb's Prometheus exporter as a collection
+// of variables with the same name and the "le" ("less than or equal") label,
+// one per bin. For example, with Bins=[1,2,10], the Prometheus variables will
+// be:
+//    myvar_here{le="1"}	12
+//    myvar_here{le="2"}	34
+//    myvar_here{le="10"}	56
+//    myvar_here{le="inf"}	78
+//
+// Additionally, a "_max", "_min" and "_count" variable will be added
+// containing the observed maximum, minimum, and total count of samples:
+//    myvar_here_max	99
+//    myvar_here_min	0
+//    myvar_here_count	180
+type Distribution struct {
+	expvar.Map
+	Bins []float64
+}
+
+func (d *Distribution) Init() {
+	// Initialze all values to zero
+	for _, bin := range d.Bins {
+		d.Map.Add(fmt.Sprint(bin), 0)
+	}
+	d.Map.Add("Inf", 0)
+	d.Map.Add("count", 0)
+	d.Map.AddFloat("min", 0.0)
+	d.Map.AddFloat("max", 0.0)
+}
+
+func (d *Distribution) AddFloat(val float64) {
+	label := "Inf"
+	for _, bin := range d.Bins {
+		if val <= bin {
+			label = fmt.Sprint(bin)
+			break
+		}
+	}
+
+	d.Map.Add(label, 1)
+	d.Map.Add("count", 1)
+
+	min, ok := d.Map.Get("min").(*expvar.Float)
+	if ok {
+		if min.Value() > val {
+			min.Set(val)
+		}
+	} else {
+		min = new(expvar.Float)
+		min.Set(val)
+		d.Map.Set("min", min)
+	}
+
+	max, ok := d.Map.Get("max").(*expvar.Float)
+	if ok {
+		if max.Value() < val {
+			max.Set(val)
+		}
+	} else {
+		max = new(expvar.Float)
+		max.Set(val)
+		d.Map.Set("max", max)
+	}
+}
diff --git a/metrics/metrics_test.go b/metrics/metrics_test.go
index 6b9aa4366..376d05c82 100644
--- a/metrics/metrics_test.go
+++ b/metrics/metrics_test.go
@@ -5,6 +5,7 @@
 package metrics
 
 import (
+	"expvar"
 	"os"
 	"runtime"
 	"testing"
@@ -51,3 +52,31 @@ func BenchmarkCurrentFileDescriptors(b *testing.B) {
 		_ = CurrentFDs()
 	}
 }
+
+func TestDistribution(t *testing.T) {
+	d := &Distribution{
+		Map: expvar.Map{},
+		Bins: []float64{
+			2, 3, 5, 8, 13,
+		},
+	}
+
+	t.Run("Single", func(t *testing.T) {
+		d.AddFloat(1.0)
+		const expected = `{"2": 1, "count": 1, "max": 1, "min": 1}`
+		if ss := d.String(); ss != expected {
+			t.Errorf("got %q; want %q", ss, expected)
+		}
+	})
+
+	t.Run("Additional", func(t *testing.T) {
+		d.AddFloat(1.5)
+		d.AddFloat(2.5)
+		d.AddFloat(7)
+		d.AddFloat(15)
+		const expected = `{"2": 2, "3": 1, "8": 1, "count": 5, "inf": 1, "max": 15, "min": 1}`
+		if ss := d.String(); ss != expected {
+			t.Errorf("got %q; want %q", ss, expected)
+		}
+	})
+}
diff --git a/tsweb/tsweb.go b/tsweb/tsweb.go
index 8ba1a52a7..2acde505c 100644
--- a/tsweb/tsweb.go
+++ b/tsweb/tsweb.go
@@ -497,6 +497,48 @@ func writePromExpVar(w io.Writer, prefix string, kv expvar.KeyValue) {
 			writePromExpVar(w, name+"_", kv)
 		})
 		return
+	case *metrics.Distribution:
+		type bucket struct {
+			le    float64
+			leStr string
+			value expvar.Var
+		}
+
+		var (
+			min, max, count expvar.Var
+			buckets         []bucket
+		)
+		v.Do(func(kv expvar.KeyValue) {
+			switch kv.Key {
+			case "min":
+				min = kv.Value
+			case "max":
+				max = kv.Value
+			case "count":
+				count = kv.Value
+			default:
+				ff, err := strconv.ParseFloat(kv.Key, 64)
+				if err == nil {
+					buckets = append(buckets, bucket{ff, kv.Key, kv.Value})
+				}
+			}
+		})
+
+		// Sort buckets by their numeric value, not string value.
+		sort.Slice(buckets, func(i, j int) bool {
+			return buckets[i].le < buckets[j].le
+		})
+
+		fmt.Fprintf(w, "# TYPE %s counter\n", name)
+		for _, bucket := range buckets {
+			fmt.Fprintf(w, "%s{le=%q} %v\n", name, bucket.leStr, bucket.value)
+		}
+
+		fmt.Fprintf(w, "# TYPE %s_min gauge\n%s_min %v\n", name, name, min)
+		fmt.Fprintf(w, "# TYPE %s_max gauge\n%s_max %v\n", name, name, max)
+		fmt.Fprintf(w, "# TYPE %s_count gauge\n%s_count %v\n", name, name, count)
+		return
+
 	case PrometheusMetricsReflectRooter:
 		root := v.PrometheusMetricsReflectRoot()
 		rv := reflect.ValueOf(root)
@@ -588,6 +630,9 @@ func writePromExpVar(w io.Writer, prefix string, kv expvar.KeyValue) {
 				fmt.Fprintf(w, "%s_%s %v\n", name, kv.Key, kv.Value)
 			})
 		}
+
+	case *metrics.Distribution:
+		// TODO
 	}
 }
 
diff --git a/tsweb/tsweb_test.go b/tsweb/tsweb_test.go
index a38b671c9..a824ed9ed 100644
--- a/tsweb/tsweb_test.go
+++ b/tsweb/tsweb_test.go
@@ -437,6 +437,56 @@ func TestVarzHandler(t *testing.T) {
 			"api_status_code_2xx 100\napi_status_code_5xx 2\n",
 		},
 		{
+			"metrics_distribution",
+			"distribution_rtt",
+			func() *metrics.Distribution {
+				d := &metrics.Distribution{
+					Bins: []float64{1, 2, 5, 10},
+				}
+				d.AddFloat(0.5)
+				d.AddFloat(4)
+				d.AddFloat(15)
+				return d
+			}(),
+			strings.TrimSpace(`
+# TYPE distribution_rtt counter
+distribution_rtt{le="1"} 1
+distribution_rtt{le="5"} 1
+distribution_rtt{le="Inf"} 1
+# TYPE distribution_rtt_min gauge
+distribution_rtt_min 0.5
+# TYPE distribution_rtt_max gauge
+distribution_rtt_max 15
+# TYPE distribution_rtt_count gauge
+distribution_rtt_count 3
+`) + "\n",
+		},
+		{
+			"metrics_distribution_empty",
+			"distribution_empty",
+			func() *metrics.Distribution {
+				d := &metrics.Distribution{
+					Bins: []float64{1, 2, 5, 10},
+				}
+				d.Init()
+				return d
+			}(),
+			strings.TrimSpace(`
+# TYPE distribution_empty counter
+distribution_empty{le="1"} 0
+distribution_empty{le="2"} 0
+distribution_empty{le="5"} 0
+distribution_empty{le="10"} 0
+distribution_empty{le="Inf"} 0
+# TYPE distribution_empty_min gauge
+distribution_empty_min 0
+# TYPE distribution_empty_max gauge
+distribution_empty_max 0
+# TYPE distribution_empty_count gauge
+distribution_empty_count 0
+`) + "\n",
+		},
+		{
 			"func_float64",
 			"counter_x",
 			expvar.Func(func() any { return float64(1.2) }),