From 68b12a74ed2379f85fa30ffe10b974cf6062c8fc Mon Sep 17 00:00:00 2001
From: James Tucker <james@tailscale.com>
Date: Thu, 19 Dec 2024 16:32:40 -0800
Subject: [PATCH] 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>
---
 cmd/stund/depaware.txt  |  2 ++
 metrics/metrics.go      | 35 +++++++++++++++++++++++++++++++++++
 metrics/metrics_test.go |  9 +++++++++
 3 files changed, 46 insertions(+)

diff --git a/cmd/stund/depaware.txt b/cmd/stund/depaware.txt
index 34a71c43e..9599f6a01 100644
--- a/cmd/stund/depaware.txt
+++ b/cmd/stund/depaware.txt
@@ -55,6 +55,7 @@ tailscale.com/cmd/stund dependencies: (generated by github.com/tailscale/depawar
         tailscale.com/net/stun                                       from tailscale.com/net/stunserver
         tailscale.com/net/stunserver                                 from tailscale.com/cmd/stund
         tailscale.com/net/tsaddr                                     from tailscale.com/tsweb
+        tailscale.com/syncs                                          from tailscale.com/metrics
         tailscale.com/tailcfg                                        from tailscale.com/version
         tailscale.com/tsweb                                          from tailscale.com/cmd/stund
         tailscale.com/tsweb/promvarz                                 from tailscale.com/tsweb
@@ -74,6 +75,7 @@ tailscale.com/cmd/stund dependencies: (generated by github.com/tailscale/depawar
    L 💣 tailscale.com/util/dirwalk                                   from tailscale.com/metrics
         tailscale.com/util/dnsname                                   from tailscale.com/tailcfg
         tailscale.com/util/lineiter                                  from tailscale.com/version/distro
+        tailscale.com/util/mak                                       from tailscale.com/syncs
         tailscale.com/util/nocasemaps                                from tailscale.com/types/ipproto
         tailscale.com/util/rands                                     from tailscale.com/tsweb
         tailscale.com/util/slicesx                                   from tailscale.com/tailcfg
diff --git a/metrics/metrics.go b/metrics/metrics.go
index a07ddccae..d1b1c06c9 100644
--- a/metrics/metrics.go
+++ b/metrics/metrics.go
@@ -11,6 +11,9 @@
 	"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
diff --git a/metrics/metrics_test.go b/metrics/metrics_test.go
index 45bf39e56..a808d5a73 100644
--- a/metrics/metrics_test.go
+++ b/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) {