net/connstats: prepare to remove package (#17554)

The connstats package was an unnecessary layer of indirection.
It was seperated out of wgengine/netlog so that net/tstun and
wgengine/magicsock wouldn't need a depenedency on the concrete
implementation of network flow logging.

Instead, we simply register a callback for counting connections.
This PR does the bare minimum work to prepare tstun and magicsock
to only care about that callback.

A future PR will delete connstats and merge it into netlog.

Updates tailscale/corp#33352

Signed-off-by: Joe Tsai <joetsai@digital-static.net>

net/tstun/wrap.go

@@ -24,7 +24,6 @@ import (
 	"go4.org/mem"
 	"tailscale.com/disco"
 	"tailscale.com/feature/buildfeatures"
-	"tailscale.com/net/connstats"
 	"tailscale.com/net/packet"
 	"tailscale.com/net/packet/checksum"
 	"tailscale.com/net/tsaddr"
@@ -33,6 +32,7 @@ import (
 	"tailscale.com/types/ipproto"
 	"tailscale.com/types/key"
 	"tailscale.com/types/logger"
+	"tailscale.com/types/netlogfunc"
 	"tailscale.com/util/clientmetric"
 	"tailscale.com/util/usermetric"
 	"tailscale.com/wgengine/filter"
@@ -203,8 +203,8 @@ type Wrapper struct {
 	// disableTSMPRejected disables TSMP rejected responses. For tests.
 	disableTSMPRejected bool
 
-	// stats maintains per-connection counters.
-	stats atomic.Pointer[connstats.Statistics]
+	// connCounter maintains per-connection counters.
+	connCounter syncs.AtomicValue[netlogfunc.ConnectionCounter]
 
 	captureHook syncs.AtomicValue[packet.CaptureCallback]
 
@@ -977,8 +977,8 @@ func (t *Wrapper) Read(buffs [][]byte, sizes []int, offset int) (int, error) {
 		}
 		sizes[buffsPos] = n
 		if buildfeatures.HasConnStats {
-			if stats := t.stats.Load(); stats != nil {
-				stats.UpdateTxVirtual(p.Buffer())
+			if update := t.connCounter.Load(); update != nil {
+				updateConnCounter(update, p.Buffer(), false)
 			}
 		}
 		buffsPos++
@@ -1106,9 +1106,9 @@ func (t *Wrapper) injectedRead(res tunInjectedRead, outBuffs [][]byte, sizes []i
 	}
 
 	if buildfeatures.HasConnStats {
-		if stats := t.stats.Load(); stats != nil {
+		if update := t.connCounter.Load(); update != nil {
 			for i := 0; i < n; i++ {
-				stats.UpdateTxVirtual(outBuffs[i][offset : offset+sizes[i]])
+				updateConnCounter(update, outBuffs[i][offset:offset+sizes[i]], false)
 			}
 		}
 	}
@@ -1276,9 +1276,9 @@ func (t *Wrapper) Write(buffs [][]byte, offset int) (int, error) {
 
 func (t *Wrapper) tdevWrite(buffs [][]byte, offset int) (int, error) {
 	if buildfeatures.HasConnStats {
-		if stats := t.stats.Load(); stats != nil {
+		if update := t.connCounter.Load(); update != nil {
 			for i := range buffs {
-				stats.UpdateRxVirtual((buffs)[i][offset:])
+				updateConnCounter(update, buffs[i][offset:], true)
 			}
 		}
 	}
@@ -1498,11 +1498,11 @@ func (t *Wrapper) Unwrap() tun.Device {
 	return t.tdev
 }
 
-// SetStatistics specifies a per-connection statistics aggregator.
+// SetConnectionCounter specifies a per-connection statistics aggregator.
 // Nil may be specified to disable statistics gathering.
-func (t *Wrapper) SetStatistics(stats *connstats.Statistics) {
+func (t *Wrapper) SetConnectionCounter(fn netlogfunc.ConnectionCounter) {
 	if buildfeatures.HasConnStats {
-		t.stats.Store(stats)
+		t.connCounter.Store(fn)
 	}
 }
 
@@ -1524,3 +1524,13 @@ func (t *Wrapper) InstallCaptureHook(cb packet.CaptureCallback) {
 	}
 	t.captureHook.Store(cb)
 }
+
+func updateConnCounter(update netlogfunc.ConnectionCounter, b []byte, receive bool) {
+	var p packet.Parsed
+	p.Decode(b)
+	if receive {
+		update(p.IPProto, p.Dst, p.Src, 1, len(b), true)
+	} else {
+		update(p.IPProto, p.Src, p.Dst, 1, len(b), false)
+	}
+}

net/tstun/wrap_test.go

@@ -5,7 +5,6 @@ package tstun
 
 import (
 	"bytes"
-	"context"
 	"encoding/binary"
 	"encoding/hex"
 	"expvar"
@@ -27,7 +26,6 @@ import (
 	"gvisor.dev/gvisor/pkg/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 	"tailscale.com/disco"
-	"tailscale.com/net/connstats"
 	"tailscale.com/net/netaddr"
 	"tailscale.com/net/packet"
 	"tailscale.com/tstest"
@@ -370,9 +368,8 @@ func TestFilter(t *testing.T) {
 	}()
 
 	var buf [MaxPacketSize]byte
-	stats := connstats.NewStatistics(0, 0, nil)
-	defer stats.Shutdown(context.Background())
-	tun.SetStatistics(stats)
+	var stats netlogtype.CountsByConnection
+	tun.SetConnectionCounter(stats.Add)
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			var n int
@@ -380,9 +377,10 @@ func TestFilter(t *testing.T) {
 			var filtered bool
 			sizes := make([]int, 1)
 
-			tunStats, _ := stats.TestExtract()
+			tunStats := stats.Clone()
+			stats.Reset()
 			if len(tunStats) > 0 {
-				t.Errorf("connstats.Statistics.Extract = %v, want {}", stats)
+				t.Errorf("connstats.Statistics.Extract = %v, want {}", tunStats)
 			}
 
 			if tt.dir == in {
@@ -415,7 +413,8 @@ func TestFilter(t *testing.T) {
 				}
 			}
 
-			got, _ := stats.TestExtract()
+			got := stats.Clone()
+			stats.Reset()
 			want := map[netlogtype.Connection]netlogtype.Counts{}
 			var wasUDP bool
 			if !tt.drop {