From 3304819739043f3def4502e345532e8e13f5e67b Mon Sep 17 00:00:00 2001
From: Maisem Ali <maisem@tailscale.com>
Date: Mon, 8 May 2023 01:22:23 -0700
Subject: [PATCH] metrics: add histogram support

Add initial histogram support.

Updates tailscale/corp#8641

Signed-off-by: Maisem Ali <maisem@tailscale.com>
---
 metrics/metrics.go | 98 +++++++++++++++++++++++++++++++++++++++++++++-
 tsweb/varz/varz.go |  2 +
 2 files changed, 99 insertions(+), 1 deletion(-)

diff --git a/metrics/metrics.go b/metrics/metrics.go
index b7339d02b..76fc9ee07 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -5,7 +5,14 @@
 // Tailscale for monitoring.
 package metrics
 
-import "expvar"
+import (
+	"expvar"
+	"fmt"
+	"io"
+	"strings"
+
+	"golang.org/x/exp/slices"
+)
 
 // Set is a string-to-Var map variable that satisfies the expvar.Var
 // interface.
@@ -66,3 +73,92 @@ func (m *LabelMap) GetFloat(key string) *expvar.Float {
 func CurrentFDs() int {
 	return currentFDs()
 }
+
+// Histogram is a histogram of values.
+// It should be created with NewHistogram.
+type Histogram struct {
+	// buckets is a list of bucket boundaries, in increasing order.
+	buckets []float64
+
+	// bucketStrings is a list of the same buckets, but as strings.
+	// This are allocated once at creation time by NewHistogram.
+	bucketStrings []string
+
+	bucketVars []expvar.Int
+	sum        expvar.Float
+	count      expvar.Int
+}
+
+// NewHistogram returns a new histogram that reports to the given
+// expvar map under the given name.
+//
+// The buckets are the boundaries of the histogram buckets, in
+// increasing order. The last bucket is +Inf.
+func NewHistogram(buckets []float64) *Histogram {
+	if !slices.IsSorted(buckets) {
+		panic("buckets must be sorted")
+	}
+	labels := make([]string, len(buckets))
+	for i, b := range buckets {
+		labels[i] = fmt.Sprintf("%v", b)
+	}
+	h := &Histogram{
+		buckets:       buckets,
+		bucketStrings: labels,
+		bucketVars:    make([]expvar.Int, len(buckets)),
+	}
+	return h
+}
+
+// Observe records a new observation in the histogram.
+func (h *Histogram) Observe(v float64) {
+	h.sum.Add(v)
+	h.count.Add(1)
+	for i, b := range h.buckets {
+		if v <= b {
+			h.bucketVars[i].Add(1)
+		}
+	}
+}
+
+// String returns a JSON representation of the histogram.
+// This is used to satisfy the expvar.Var interface.
+func (h *Histogram) String() string {
+	var b strings.Builder
+	fmt.Fprintf(&b, "{")
+	first := true
+	h.Do(func(kv expvar.KeyValue) {
+		if !first {
+			fmt.Fprintf(&b, ",")
+		}
+		fmt.Fprintf(&b, "%q: ", kv.Key)
+		if kv.Value != nil {
+			fmt.Fprintf(&b, "%v", kv.Value)
+		} else {
+			fmt.Fprint(&b, "null")
+		}
+		first = false
+	})
+	fmt.Fprintf(&b, "\"sum\": %v,", &h.sum)
+	fmt.Fprintf(&b, "\"count\": %v", &h.count)
+	fmt.Fprintf(&b, "}")
+	return b.String()
+}
+
+// Do calls f for each bucket in the histogram.
+func (h *Histogram) Do(f func(expvar.KeyValue)) {
+	for i := range h.bucketVars {
+		f(expvar.KeyValue{Key: h.bucketStrings[i], Value: &h.bucketVars[i]})
+	}
+	f(expvar.KeyValue{Key: "+Inf", Value: &h.count})
+}
+
+// PromExport writes the histogram to w in Prometheus exposition format.
+func (h *Histogram) PromExport(w io.Writer, name string) {
+	fmt.Fprintf(w, "# TYPE %s histogram\n", name)
+	h.Do(func(kv expvar.KeyValue) {
+		fmt.Fprintf(w, "%s_bucket{le=%q} %v\n", name, kv.Key, kv.Value)
+	})
+	fmt.Fprintf(w, "%s_sum %v\n", name, &h.sum)
+	fmt.Fprintf(w, "%s_count %v\n", name, &h.count)
+}
diff --git a/tsweb/varz/varz.go b/tsweb/varz/varz.go
index 12922106e..858bbcc3d 100644
--- a/tsweb/varz/varz.go
+++ b/tsweb/varz/varz.go
@@ -191,6 +191,8 @@ func writePromExpVar(w io.Writer, prefix string, kv expvar.KeyValue) {
 		v.Do(func(kv expvar.KeyValue) {
 			fmt.Fprintf(w, "%s{%s=%q} %v\n", name, cmpx.Or(v.Label, "label"), kv.Key, kv.Value)
 		})
+	case *metrics.Histogram:
+		v.PromExport(w, name)
 	case *expvar.Map:
 		if label != "" && typ != "" {
 			fmt.Fprintf(w, "# TYPE %s %s\n", name, typ)