metrics,syncs: add ShardedInt support to metrics.LabelMap

metrics.LabelMap grows slightly more heavy, needing a lock to ensure
proper ordering for newly initialized ShardedInt values. An Add method
enables callers to use .Add for both expvar.Int and syncs.ShardedInt
values, but retains the original behavior of defaulting to initializing
expvar.Int values.

Updates tailscale/corp#25450

Co-Authored-By: Andrew Dunham <andrew@du.nham.ca>
Signed-off-by: James Tucker <james@tailscale.com>

metrics/metrics.go

@@ -11,6 +11,9 @@ import (
 	"io"
 	"slices"
 	"strings"
+	"sync"
+
+	"tailscale.com/syncs"
 )
 
 // Set is a string-to-Var map variable that satisfies the expvar.Var
@@ -37,6 +40,8 @@ type Set struct {
 type LabelMap struct {
 	Label string
 	expvar.Map
+	// shardedIntMu orders the initialization of new shardedint keys
+	shardedIntMu sync.Mutex
 }
 
 // SetInt64 sets the *Int value stored under the given map key.
@@ -44,6 +49,19 @@ func (m *LabelMap) SetInt64(key string, v int64) {
 	m.Get(key).Set(v)
 }
 
+// Add adds delta to the any int-like value stored under the given map key.
+func (m *LabelMap) Add(key string, delta int64) {
+	type intAdder interface {
+		Add(delta int64)
+	}
+	o := m.Map.Get(key)
+	if o == nil {
+		m.Map.Add(key, delta)
+		return
+	}
+	o.(intAdder).Add(delta)
+}
+
 // Get returns a direct pointer to the expvar.Int for key, creating it
 // if necessary.
 func (m *LabelMap) Get(key string) *expvar.Int {
@@ -51,6 +69,23 @@ func (m *LabelMap) Get(key string) *expvar.Int {
 	return m.Map.Get(key).(*expvar.Int)
 }
 
+// GetShardedInt returns a direct pointer to the syncs.ShardedInt for key,
+// creating it if necessary.
+func (m *LabelMap) GetShardedInt(key string) *syncs.ShardedInt {
+	i := m.Map.Get(key)
+	if i == nil {
+		m.shardedIntMu.Lock()
+		defer m.shardedIntMu.Unlock()
+		i = m.Map.Get(key)
+		if i != nil {
+			return i.(*syncs.ShardedInt)
+		}
+		i = syncs.NewShardedInt()
+		m.Set(key, i)
+	}
+	return i.(*syncs.ShardedInt)
+}
+
 // GetIncrFunc returns a function that increments the expvar.Int named by key.
 //
 // Most callers should not need this; it exists to satisfy an

metrics/metrics_test.go

@@ -21,6 +21,15 @@ func TestLabelMap(t *testing.T) {
 	if g, w := m.Get("bar").Value(), int64(2); g != w {
 		t.Errorf("bar = %v; want %v", g, w)
 	}
+	m.GetShardedInt("sharded").Add(5)
+	if g, w := m.GetShardedInt("sharded").Value(), int64(5); g != w {
+		t.Errorf("sharded = %v; want %v", g, w)
+	}
+	m.Add("sharded", 1)
+	if g, w := m.GetShardedInt("sharded").Value(), int64(6); g != w {
+		t.Errorf("sharded = %v; want %v", g, w)
+	}
+	m.Add("neverbefore", 1)
 }
 
 func TestCurrentFileDescriptors(t *testing.T) {