sockstats: switch label to enum

Makes it cheaper/simpler to persist values, and encourages reuse of
labels as opposed to generating an arbitrary number.

Updates tailscale/corp#9230
Updates #3363

Signed-off-by: Mihai Parparita <[email protected]>

control/controlclient/auto.go

@@ -119,10 +119,10 @@ func NewNoStart(opts Options) (_ *Auto, err error) {
 		statusFunc: opts.Status,
 	}
 	c.authCtx, c.authCancel = context.WithCancel(context.Background())
-	c.authCtx = sockstats.WithSockStats(c.authCtx, "controlclient.Auto:auth")
+	c.authCtx = sockstats.WithSockStats(c.authCtx, sockstats.LabelControlClientAuto)
 
 	c.mapCtx, c.mapCancel = context.WithCancel(context.Background())
-	c.mapCtx = sockstats.WithSockStats(c.mapCtx, "controlclient.Auto:map")
+	c.mapCtx = sockstats.WithSockStats(c.mapCtx, sockstats.LabelControlClientAuto)
 
 	c.unregisterHealthWatch = health.RegisterWatcher(direct.ReportHealthChange)
 	return c, nil
@@ -211,7 +211,7 @@ func (c *Auto) cancelAuth() {
 	}
 	if !c.closed {
 		c.authCtx, c.authCancel = context.WithCancel(context.Background())
-		c.authCtx = sockstats.WithSockStats(c.authCtx, "controlclient.Auto:auth")
+		c.authCtx = sockstats.WithSockStats(c.authCtx, sockstats.LabelControlClientAuto)
 	}
 	c.mu.Unlock()
 }
@@ -222,7 +222,7 @@ func (c *Auto) cancelMapLocked() {
 	}
 	if !c.closed {
 		c.mapCtx, c.mapCancel = context.WithCancel(context.Background())
-		c.mapCtx = sockstats.WithSockStats(c.mapCtx, "controlclient.Auto:map")
+		c.mapCtx = sockstats.WithSockStats(c.mapCtx, sockstats.LabelControlClientAuto)
 
 	}
 }

control/controlhttp/client.go

@@ -273,7 +273,7 @@ func (a *Dialer) dialHost(ctx context.Context, addr netip.Addr) (*ClientConn, er
 	ctx, cancel := context.WithCancel(ctx)
 	defer cancel()
 
-	ctx = sockstats.WithSockStats(ctx, "controlclient.Dialer")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelControlClientDialer)
 
 	// u80 and u443 are the URLs we'll try to hit over HTTP or HTTPS,
 	// respectively, in order to do the HTTP upgrade to a net.Conn over which

derp/derphttp/derphttp_client.go

@@ -616,7 +616,7 @@ func (c *Client) dialNode(ctx context.Context, n *tailcfg.DERPNode) (net.Conn, e
 	ctx, cancel := context.WithTimeout(ctx, dialNodeTimeout)
 	defer cancel()
 
-	ctx = sockstats.WithSockStats(ctx, "derphttp.Client")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelDERPHTTPClient)
 
 	nwait := 0
 	startDial := func(dstPrimary, proto string) {

ipn/ipnlocal/peerapi.go

@@ -879,11 +879,13 @@ func (h *peerAPIHandler) handleServeSockStats(w http.ResponseWriter, r *http.Req
 	fmt.Fprintln(w, "</thead>")
 
 	fmt.Fprintln(w, "<tbody>")
-	labels := make([]string, 0, len(stats.Stats))
+	labels := make([]sockstats.Label, 0, len(stats.Stats))
 	for label := range stats.Stats {
 		labels = append(labels, label)
 	}
-	sort.Strings(labels)
+	slices.SortFunc(labels, func(a, b sockstats.Label) bool {
+		return a.String() < b.String()
+	})
 
 	txTotal := int64(0)
 	rxTotal := int64(0)
@@ -893,7 +895,7 @@ func (h *peerAPIHandler) handleServeSockStats(w http.ResponseWriter, r *http.Req
 	for _, label := range labels {
 		stat := stats.Stats[label]
 		fmt.Fprintln(w, "<tr>")
-		fmt.Fprintf(w, "<td>%s</td>", html.EscapeString(label))
+		fmt.Fprintf(w, "<td>%s</td>", html.EscapeString(label.String()))
 		fmt.Fprintf(w, "<td align=right>%d</td>", stat.TxBytes)
 		fmt.Fprintf(w, "<td align=right>%d</td>", stat.RxBytes)
 

logtail/logtail.go

@@ -428,7 +428,7 @@ func (l *Logger) awaitInternetUp(ctx context.Context) {
 // origlen of -1 indicates that the body is not compressed.
 func (l *Logger) upload(ctx context.Context, body []byte, origlen int) (uploaded bool, err error) {
 	const maxUploadTime = 45 * time.Second
-	ctx = sockstats.WithSockStats(ctx, "logtail.Logger")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelLogtailLogger)
 	ctx, cancel := context.WithTimeout(ctx, maxUploadTime)
 	defer cancel()
 

net/dns/resolver/forwarder.go

@@ -407,7 +407,7 @@ func (f *forwarder) getKnownDoHClientForProvider(urlBase string) (c *http.Client
 const dohType = "application/dns-message"
 
 func (f *forwarder) sendDoH(ctx context.Context, urlBase string, c *http.Client, packet []byte) ([]byte, error) {
-	ctx = sockstats.WithSockStats(ctx, "dns.forwarder:doh")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelDNSForwarderDoH)
 	metricDNSFwdDoH.Add(1)
 	req, err := http.NewRequestWithContext(ctx, "POST", urlBase, bytes.NewReader(packet))
 	if err != nil {
@@ -487,7 +487,7 @@ func (f *forwarder) sendUDP(ctx context.Context, fq *forwardQuery, rr resolverAn
 		return nil, fmt.Errorf("unrecognized resolver type %q", rr.name.Addr)
 	}
 	metricDNSFwdUDP.Add(1)
-	ctx = sockstats.WithSockStats(ctx, "dns.forwarder:udp")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelDNSForwarderUDP)
 
 	ln, err := f.packetListener(ipp.Addr())
 	if err != nil {

net/netcheck/netcheck.go

@@ -784,7 +784,7 @@ func (c *Client) GetReport(ctx context.Context, dm *tailcfg.DERPMap) (_ *Report,
 	ctx, cancel := context.WithTimeout(ctx, overallProbeTimeout)
 	defer cancel()
 
-	ctx = sockstats.WithSockStats(ctx, "netcheck.Client")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelNetcheckClient)
 
 	if dm == nil {
 		return nil, errors.New("netcheck: GetReport: DERP map is nil")

net/portmapper/portmapper.go

@@ -249,7 +249,7 @@ func (c *Client) upnpPort() uint16 {
 }
 
 func (c *Client) listenPacket(ctx context.Context, network, addr string) (nettype.PacketConn, error) {
-	ctx = sockstats.WithSockStats(ctx, "portmapper.Client")
+	ctx = sockstats.WithSockStats(ctx, sockstats.LabelPortmapperClient)
 
 	// When running under testing conditions, we bind the IGD server
 	// to localhost, and may be running in an environment where our

net/sockstats/label_string.go

@@ -0,0 +1,32 @@
+// Code generated by "stringer -type Label -trimprefix Label"; DO NOT EDIT.
+
+package sockstats
+
+import "strconv"
+
+func _() {
+	// An "invalid array index" compiler error signifies that the constant values have changed.
+	// Re-run the stringer command to generate them again.
+	var x [1]struct{}
+	_ = x[LabelControlClientAuto-0]
+	_ = x[LabelControlClientDialer-1]
+	_ = x[LabelDERPHTTPClient-2]
+	_ = x[LabelLogtailLogger-3]
+	_ = x[LabelDNSForwarderDoH-4]
+	_ = x[LabelDNSForwarderUDP-5]
+	_ = x[LabelNetcheckClient-6]
+	_ = x[LabelPortmapperClient-7]
+	_ = x[LabelMagicsockConnUDP4-8]
+	_ = x[LabelMagicsockConnUDP6-9]
+}
+
+const _Label_name = "ControlClientAutoControlClientDialerDERPHTTPClientLogtailLoggerDNSForwarderDoHDNSForwarderUDPNetcheckClientPortmapperClientMagicsockConnUDP4MagicsockConnUDP6"
+
+var _Label_index = [...]uint8{0, 17, 36, 50, 63, 78, 93, 107, 123, 140, 157}
+
+func (i Label) String() string {
+	if i >= Label(len(_Label_index)-1) {
+		return "Label(" + strconv.FormatInt(int64(i), 10) + ")"
+	}
+	return _Label_name[_Label_index[i]:_Label_index[i+1]]
+}

net/sockstats/sockstats.go

@@ -15,10 +15,32 @@ import (
 )
 
 type SockStats struct {
-	Stats      map[string]SockStat
+	Stats      map[Label]SockStat
 	Interfaces []string
 }
 
+// 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.
+type Label uint8
+
+//go:generate go run golang.org/x/tools/cmd/stringer -type Label -trimprefix Label
+
+// Labels are named after the package and function/struct that uses the socket.
+// Values may be persisted and thus existing entries should not be re-numbered.
+const (
+	LabelControlClientAuto   Label = 0 // control/controlclient/auto.go
+	LabelControlClientDialer Label = 1 // control/controlhttp/client.go
+	LabelDERPHTTPClient      Label = 2 // derp/derphttp/derphttp_client.go
+	LabelLogtailLogger       Label = 3 // logtail/logtail.go
+	LabelDNSForwarderDoH     Label = 4 // net/dns/resolver/forwarder.go
+	LabelDNSForwarderUDP     Label = 5 // net/dns/resolver/forwarder.go
+	LabelNetcheckClient      Label = 6 // net/netcheck/netcheck.go
+	LabelPortmapperClient    Label = 7 // net/portmapper/portmapper.go
+	LabelMagicsockConnUDP4   Label = 8 // wgengine/magicsock/magicsock.go
+	LabelMagicsockConnUDP6   Label = 9 // wgengine/magicsock/magicsock.go
+)
+
 type SockStat struct {
 	TxBytes            int64
 	RxBytes            int64
@@ -26,7 +48,7 @@ type SockStat struct {
 	RxBytesByInterface map[string]int64
 }
 
-func WithSockStats(ctx context.Context, label string) context.Context {
+func WithSockStats(ctx context.Context, label Label) context.Context {
 	return withSockStats(ctx, label)
 }
 

net/sockstats/sockstats_noop.go

@@ -9,7 +9,7 @@ import (
 	"context"
 )
 
-func withSockStats(ctx context.Context, label string) context.Context {
+func withSockStats(ctx context.Context, label Label) context.Context {
 	return ctx
 }
 

net/sockstats/sockstats_tsgo.go

@@ -24,7 +24,7 @@ var sockStats = struct {
 	// mu protects fields in this group. It should not be held in the per-read/
 	// write callbacks.
 	mu              sync.Mutex
-	countersByLabel map[string]*sockStatCounters
+	countersByLabel map[Label]*sockStatCounters
 	knownInterfaces map[int]string // interface index -> name
 	usedInterfaces  map[int]int    // set of interface indexes
 
@@ -32,12 +32,12 @@ var sockStats = struct {
 	// write callbacks.
 	currentInterface atomic.Uint32
 }{
-	countersByLabel: make(map[string]*sockStatCounters),
+	countersByLabel: make(map[Label]*sockStatCounters),
 	knownInterfaces: make(map[int]string),
 	usedInterfaces:  make(map[int]int),
 }
 
-func withSockStats(ctx context.Context, label string) context.Context {
+func withSockStats(ctx context.Context, label Label) context.Context {
 	sockStats.mu.Lock()
 	defer sockStats.mu.Unlock()
 	counters, ok := sockStats.countersByLabel[label]
@@ -85,7 +85,7 @@ func get() *SockStats {
 	defer sockStats.mu.Unlock()
 
 	r := &SockStats{
-		Stats:      make(map[string]SockStat),
+		Stats:      make(map[Label]SockStat),
 		Interfaces: make([]string, 0, len(sockStats.usedInterfaces)),
 	}
 	for iface := range sockStats.usedInterfaces {

wgengine/magicsock/magicsock.go

@@ -3050,7 +3050,12 @@ func (c *Conn) ReSTUN(why string) {
 // listenPacket opens a packet listener.
 // The network must be "udp4" or "udp6".
 func (c *Conn) listenPacket(network string, port uint16) (nettype.PacketConn, error) {
-	ctx := sockstats.WithSockStats(context.Background(), fmt.Sprintf("magicsock.Conn:%s", network)) // unused without DNS name to resolve
+	ctx := context.Background() // unused without DNS name to resolve
+	if network == "udp4" {
+		ctx = sockstats.WithSockStats(ctx, sockstats.LabelMagicsockConnUDP4)
+	} else {
+		ctx = sockstats.WithSockStats(ctx, sockstats.LabelMagicsockConnUDP6)
+	}
 	addr := net.JoinHostPort("", fmt.Sprint(port))
 	if c.testOnlyPacketListener != nil {
 		return nettype.MakePacketListenerWithNetIP(c.testOnlyPacketListener).ListenPacket(ctx, network, addr)