wgengine: instrument with usermetrics

Updates tailscale/corp#22075

Signed-off-by: Kristoffer Dalby <kristoffer@tailscale.com>

util/clientmetric/clientmetric.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/hex"
+	"expvar"
 	"fmt"
 	"io"
 	"sort"
@@ -16,6 +17,8 @@ import (
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"tailscale.com/util/set"
 )
 
 var (
@@ -223,6 +226,54 @@ func NewGaugeFunc(name string, f func() int64) *Metric {
 	return m
 }
 
+// AggregateCounter returns a sum of expvar counters registered with it.
+type AggregateCounter struct {
+	mu       sync.RWMutex
+	counters set.Set[*expvar.Int]
+}
+
+func (c *AggregateCounter) Value() int64 {
+	c.mu.RLock()
+	defer c.mu.RUnlock()
+	var sum int64
+	for cnt := range c.counters {
+		sum += cnt.Value()
+	}
+	return sum
+}
+
+// Register registers provided expvar counter.
+// When a counter is added to the counter, it will be reset
+// to start counting from 0. This is to avoid incrementing the
+// counter with an unexpectedly large value.
+func (c *AggregateCounter) Register(counter *expvar.Int) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	// No need to do anything if it's already registered.
+	if c.counters.Contains(counter) {
+		return
+	}
+	counter.Set(0)
+	c.counters.Add(counter)
+}
+
+// UnregisterAll unregisters all counters resulting in it
+// starting back down at zero. This is to ensure monotonicity
+// and respect the semantics of the counter.
+func (c *AggregateCounter) UnregisterAll() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.counters = set.Set[*expvar.Int]{}
+}
+
+// NewAggregateCounter returns a new aggregate counter that returns
+// a sum of expvar variables registered with it.
+func NewAggregateCounter(name string) *AggregateCounter {
+	c := &AggregateCounter{counters: set.Set[*expvar.Int]{}}
+	NewGaugeFunc(name, c.Value)
+	return c
+}
+
 // WritePrometheusExpositionFormat writes all client metrics to w in
 // the Prometheus text-based exposition format.
 //

util/clientmetric/clientmetric_test.go

@@ -4,8 +4,11 @@
 package clientmetric
 
 import (
+	"expvar"
 	"testing"
 	"time"
+
+	qt "github.com/frankban/quicktest"
 )
 
 func TestDeltaEncBuf(t *testing.T) {
@@ -107,3 +110,49 @@ func TestWithFunc(t *testing.T) {
 		t.Errorf("second = %q; want %q", got, want)
 	}
 }
+
+func TestAggregateCounter(t *testing.T) {
+	clearMetrics()
+
+	c := qt.New(t)
+
+	expv1 := &expvar.Int{}
+	expv2 := &expvar.Int{}
+	expv3 := &expvar.Int{}
+
+	aggCounter := NewAggregateCounter("agg_counter")
+
+	aggCounter.Register(expv1)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(0))
+
+	expv1.Add(1)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(1))
+
+	aggCounter.Register(expv2)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(1))
+
+	expv1.Add(1)
+	expv2.Add(1)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(3))
+
+	// Adding a new expvar should not change the value
+	// and any value the counter already had is reset
+	expv3.Set(5)
+	aggCounter.Register(expv3)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(3))
+
+	// Registering the same expvar multiple times should not change the value
+	aggCounter.Register(expv3)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(3))
+
+	aggCounter.UnregisterAll()
+	c.Assert(aggCounter.Value(), qt.Equals, int64(0))
+
+	// Start over
+	expv3.Set(5)
+	aggCounter.Register(expv3)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(0))
+
+	expv3.Set(5)
+	c.Assert(aggCounter.Value(), qt.Equals, int64(5))
+}