net/tunstats: new package to track per-connection counters (#5818)

High-level API:

	type Statistics struct { ... }
	type Counts struct { TxPackets, TxBytes, RxPackets, RxBytes uint64 }
	func (*Statistics) UpdateTx([]byte)
	func (*Statistics) UpdateRx([]byte)
	func (*Statistics) Extract() map[flowtrack.Tuple]Counts

The API accepts a []byte instead of a packet.Parsed so that a future
implementation can directly hash the address and port bytes,
which are contiguous in most IP packets.
This will be useful for a custom concurrent-safe hashmap implementation.

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

net/tunstats/stats.go

@@ -0,0 +1,78 @@
+// Copyright (c) 2022 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package tunstats maintains statistics about connections
+// flowing through a TUN device (which operate at the IP layer).
+package tunstats
+
+import (
+	"sync"
+
+	"tailscale.com/net/flowtrack"
+	"tailscale.com/net/packet"
+)
+
+// Statistics maintains counters for every connection.
+// All methods are safe for concurrent use.
+// The zero value is ready for use.
+type Statistics struct {
+	mu sync.Mutex
+	m  map[flowtrack.Tuple]Counts
+}
+
+// Counts are statistics about a particular connection.
+type Counts struct {
+	TxPackets uint64 `json:"txPkts,omitempty"`
+	TxBytes   uint64 `json:"txBytes,omitempty"`
+	RxPackets uint64 `json:"rxPkts,omitempty"`
+	RxBytes   uint64 `json:"rxBytes,omitempty"`
+}
+
+// UpdateTx updates the counters for a transmitted IP packet
+// The source and destination of the packet directly correspond with
+// the source and destination in flowtrack.Tuple.
+func (s *Statistics) UpdateTx(b []byte) {
+	s.update(b, false)
+}
+
+// UpdateRx updates the counters for a received IP packet.
+// The source and destination of the packet are inverted with respect to
+// the source and destination in flowtrack.Tuple.
+func (s *Statistics) UpdateRx(b []byte) {
+	s.update(b, true)
+}
+
+func (s *Statistics) update(b []byte, receive bool) {
+	var p packet.Parsed
+	p.Decode(b)
+	tuple := flowtrack.Tuple{Proto: p.IPProto, Src: p.Src, Dst: p.Dst}
+	if receive {
+		tuple.Src, tuple.Dst = tuple.Dst, tuple.Src
+	}
+
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	if s.m == nil {
+		s.m = make(map[flowtrack.Tuple]Counts)
+	}
+	cnts := s.m[tuple]
+	if receive {
+		cnts.RxPackets++
+		cnts.RxBytes += uint64(len(b))
+	} else {
+		cnts.TxPackets++
+		cnts.TxBytes += uint64(len(b))
+	}
+	s.m[tuple] = cnts
+}
+
+// Extract extracts and resets the counters for all active connections.
+// It must be called periodically otherwise the memory used is unbounded.
+func (s *Statistics) Extract() map[flowtrack.Tuple]Counts {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	m := s.m
+	s.m = make(map[flowtrack.Tuple]Counts)
+	return m
+}