sockstats: refactor validation to be opt-in

Followup to #7499 to make validation a separate function (
GetWithValidation vs. Get). This way callers that don't need it don't
pay the cost of a syscall per active TCP socket.

Also clears the conn on close, so that we don't double-count the stats.

Also more consistently uses Go doc comments for the exported API of the
sockstats package.

Updates tailscale/corp#9230
Updates #3363

Signed-off-by: Mihai Parparita <mihai@tailscale.com>

net/sockstats/sockstats.go

@@ -14,11 +14,25 @@ import (
 	"tailscale.com/net/interfaces"
 )
 
+// SockStats contains statistics for sockets instrumented with the
+// WithSockStats() function, along with the interfaces that we have
+// per-interface statistics for.
 type SockStats struct {
 	Stats      map[Label]SockStat
 	Interfaces []string
 }
 
+// SockStat contains the sent and received bytes for a socket instrumented with
+// the WithSockStats() function. The bytes are also broken down by interface,
+// though this may be a subset of the total if interfaces were added after the
+// instrumented socket was created.
+type SockStat struct {
+	TxBytes            uint64
+	RxBytes            uint64
+	TxBytesByInterface map[string]uint64
+	RxBytesByInterface map[string]uint64
+}
+
 // Label is an identifier for a socket that stats are collected for. A finite
 // set of values that may be used to label a socket to encourage grouping and
 // to make storage more efficient.
@@ -41,25 +55,38 @@ const (
 	LabelMagicsockConnUDP6   Label = 9 // wgengine/magicsock/magicsock.go
 )
 
-type SockStat struct {
-	TxBytes            uint64
-	RxBytes            uint64
-	TxBytesByInterface map[string]uint64
-	RxBytesByInterface map[string]uint64
-
-	// NOCOMMIT
-	ValidationTxBytes uint64
-	ValidationRxBytes uint64
-}
-
+// WithSockStats instruments a context so that sockets created with it will
+// have their statistics collected.
 func WithSockStats(ctx context.Context, label Label) context.Context {
 	return withSockStats(ctx, label)
 }
 
+// Get returns the current socket statistics.
 func Get() *SockStats {
 	return get()
 }
 
+// ValidationSockStats contains external validation numbers for sockets
+// instrumented with WithSockStats. It may be a subset of the all sockets,
+// depending on what externa measurement mechanisms the platform supports.
+type ValidationSockStats struct {
+	Stats map[Label]ValidationSockStat
+}
+
+// ValidationSockStat contains the validation bytes for a socket instrumented
+// with WithSockStats.
+type ValidationSockStat struct {
+	TxBytes uint64
+	RxBytes uint64
+}
+
+// GetWithValidation is a variant of GetWith that returns both the current stats
+// and external validation numbers for the stats. It is more expensive than
+// Get and should be used in debug interfaces only.
+func GetWithValidation() (*SockStats, *ValidationSockStats) {
+	return get(), getValidation()
+}
+
 // LinkMonitor is the interface for the parts of wgengine/mointor's Mon that we
 // need, to avoid the dependency.
 type LinkMonitor interface {
@@ -67,6 +94,8 @@ type LinkMonitor interface {
 	RegisterChangeCallback(interfaces.ChangeFunc) (unregister func())
 }
 
+// SetLinkMonitor configures the sockstats package to monitor the active
+// interface, so that per-interface stats can be collected.
 func SetLinkMonitor(lm LinkMonitor) {
 	setLinkMonitor(lm)
 }

net/sockstats/sockstats_noop.go

@@ -17,5 +17,9 @@ func get() *SockStats {
 	return nil
 }
 
+func getValidation() *ValidationSockStats {
+	return nil
+}
+
 func setLinkMonitor(lm LinkMonitor) {
 }

net/sockstats/sockstats_tsgo.go

@@ -96,6 +96,7 @@ func withSockStats(ctx context.Context, label Label) context.Context {
 		tx, rx := tcpConnStats(c)
 		counters.validationTxBytes.Add(tx)
 		counters.validationRxBytes.Add(rx)
+		counters.validationConn.Store(nil)
 	}
 
 	// Don't bother adding these hooks if we can't get stats that they end up
@@ -167,14 +168,6 @@ func get() *SockStats {
 			RxBytes:            counters.rxBytes.Load(),
 			TxBytesByInterface: make(map[string]uint64),
 			RxBytesByInterface: make(map[string]uint64),
-
-			ValidationTxBytes: counters.validationTxBytes.Load(),
-			ValidationRxBytes: counters.validationRxBytes.Load(),
-		}
-		if c := counters.validationConn.Load(); c != nil && tcpConnStats != nil {
-			tx, rx := tcpConnStats(*c)
-			s.ValidationTxBytes += tx
-			s.ValidationRxBytes += rx
 		}
 		for iface, a := range counters.rxBytesByInterface {
 			ifName := sockStats.knownInterfaces[iface]
@@ -190,6 +183,30 @@ func get() *SockStats {
 	return r
 }
 
+func getValidation() *ValidationSockStats {
+	sockStats.mu.Lock()
+	defer sockStats.mu.Unlock()
+
+	r := &ValidationSockStats{
+		Stats: make(map[Label]ValidationSockStat),
+	}
+
+	for label, counters := range sockStats.countersByLabel {
+		s := ValidationSockStat{
+			TxBytes: counters.validationTxBytes.Load(),
+			RxBytes: counters.validationRxBytes.Load(),
+		}
+		if c := counters.validationConn.Load(); c != nil && tcpConnStats != nil {
+			tx, rx := tcpConnStats(*c)
+			s.TxBytes += tx
+			s.RxBytes += rx
+		}
+		r.Stats[label] = s
+	}
+
+	return r
+}
+
 func setLinkMonitor(lm LinkMonitor) {
 	sockStats.mu.Lock()
 	defer sockStats.mu.Unlock()