net/flowtrack: optimize Tuple type for use as map key

This gets UDP filter overhead closer to TCP. Still ~2x, but no longer ~3x.

    goos: darwin
    goarch: arm64
    pkg: tailscale.com/wgengine/filter
                                       │   before    │                after                │
                                       │   sec/op    │   sec/op     vs base                │
    FilterMatch/tcp-not-syn-v4-8         15.43n ± 3%   15.38n ± 5%        ~ (p=0.339 n=10)
    FilterMatch/udp-existing-flow-v4-8   42.45n ± 0%   34.77n ± 1%  -18.08% (p=0.000 n=10)
    geomean                              25.59n        23.12n        -9.65%

Updates #12486

Change-Id: I595cfadcc6b7234604bed9c4dd4261e087c0d4c4
Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>

net/flowtrack/flowtrack.go

@@ -11,23 +11,81 @@ package flowtrack
 
 import (
 	"container/list"
+	"encoding/json"
 	"fmt"
 	"net/netip"
 
 	"tailscale.com/types/ipproto"
 )
 
+// MakeTuple makes a Tuple out of netip.AddrPort values.
+func MakeTuple(proto ipproto.Proto, src, dst netip.AddrPort) Tuple {
+	return Tuple{
+		proto:   proto,
+		src:     src.Addr().As16(),
+		srcPort: src.Port(),
+		dst:     dst.Addr().As16(),
+		dstPort: dst.Port(),
+	}
+}
+
 // Tuple is a 5-tuple of proto, source and destination IP and port.
+//
+// This struct originally used netip.AddrPort, but that was about twice as slow
+// when used as a map key due to the alignment and extra space for the IPv6 zone
+// pointers (unneeded for all our current 2024-06-17 flowtrack needs).
+//
+// This struct is packed optimally and doesn't contain gaps or pointers.
 type Tuple struct {
+	src     [16]byte
+	dst     [16]byte
+	srcPort uint16
+	dstPort uint16
+	proto   ipproto.Proto
+}
+
+func (t Tuple) SrcAddr() netip.Addr {
+	return netip.AddrFrom16(t.src).Unmap()
+}
+
+func (t Tuple) DstAddr() netip.Addr {
+	return netip.AddrFrom16(t.dst).Unmap()
+}
+
+func (t Tuple) SrcPort() uint16 { return t.srcPort }
+func (t Tuple) DstPort() uint16 { return t.dstPort }
+
+func (t Tuple) String() string {
+	return fmt.Sprintf("(%v %v => %v)", t.proto, t.src, t.dst)
+}
+
+func (t Tuple) MarshalJSON() ([]byte, error) {
+	return json.Marshal(tupleOld{
+		Proto: t.proto,
+		Src:   netip.AddrPortFrom(t.SrcAddr(), t.srcPort),
+		Dst:   netip.AddrPortFrom(t.DstAddr(), t.dstPort),
+	})
+}
+
+func (t *Tuple) UnmarshalJSON(b []byte) error {
+	var ot tupleOld
+	if err := json.Unmarshal(b, &ot); err != nil {
+		return err
+	}
+	*t = MakeTuple(ot.Proto, ot.Src, ot.Dst)
+	return nil
+}
+
+// tupleOld is the old JSON representation of Tuple, before
+// we split and rearranged the fields for efficiency. This type
+// is the JSON adapter type to make sure we still generate
+// the same JSON as before.
+type tupleOld struct {
 	Proto ipproto.Proto  `json:"proto"`
 	Src   netip.AddrPort `json:"src"`
 	Dst   netip.AddrPort `json:"dst"`
 }
 
-func (t Tuple) String() string {
-	return fmt.Sprintf("(%v %v => %v)", t.Proto, t.Src, t.Dst)
-}
-
 // Cache is an LRU cache keyed by Tuple.
 //
 // The zero value is valid to use.

net/flowtrack/flowtrack_test.go

@@ -4,19 +4,21 @@
 package flowtrack
 
 import (
+	"encoding/json"
 	"net/netip"
 	"testing"
 
 	"tailscale.com/tstest"
+	"tailscale.com/types/ipproto"
 )
 
 func TestCache(t *testing.T) {
 	c := &Cache[int]{MaxEntries: 2}
 
-	k1 := Tuple{Src: netip.MustParseAddrPort("1.1.1.1:1"), Dst: netip.MustParseAddrPort("1.1.1.1:1")}
-	k2 := Tuple{Src: netip.MustParseAddrPort("1.1.1.1:1"), Dst: netip.MustParseAddrPort("2.2.2.2:2")}
-	k3 := Tuple{Src: netip.MustParseAddrPort("1.1.1.1:1"), Dst: netip.MustParseAddrPort("3.3.3.3:3")}
-	k4 := Tuple{Src: netip.MustParseAddrPort("1.1.1.1:1"), Dst: netip.MustParseAddrPort("4.4.4.4:4")}
+	k1 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("1.1.1.1:1"))
+	k2 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("2.2.2.2:2"))
+	k3 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("3.3.3.3:3"))
+	k4 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("4.4.4.4:4"))
 
 	wantLen := func(want int) {
 		t.Helper()
@@ -80,3 +82,42 @@ func TestCache(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func BenchmarkMapKeys(b *testing.B) {
+	b.Run("typed", func(b *testing.B) {
+		c := &Cache[struct{}]{MaxEntries: 1000}
+		var t Tuple
+		for proto := range 20 {
+			t = Tuple{proto: ipproto.Proto(proto), src: netip.MustParseAddr("1.1.1.1").As16(), srcPort: 1, dst: netip.MustParseAddr("1.1.1.1").As16(), dstPort: 1}
+			c.Add(t, struct{}{})
+		}
+		for i := 0; i < b.N; i++ {
+			_, ok := c.Get(t)
+			if !ok {
+				b.Fatal("missing key")
+			}
+		}
+	})
+}
+
+func TestJSON(t *testing.T) {
+	v := MakeTuple(123,
+		netip.MustParseAddrPort("1.2.3.4:5"),
+		netip.MustParseAddrPort("6.7.8.9:10"))
+	got, err := json.Marshal(v)
+	if err != nil {
+		t.Fatal(err)
+	}
+	const want = `{"proto":123,"src":"1.2.3.4:5","dst":"6.7.8.9:10"}`
+	if string(got) != want {
+		t.Errorf("Marshal = %q; want %q", got, want)
+	}
+
+	var back Tuple
+	if err := json.Unmarshal(got, &back); err != nil {
+		t.Fatal(err)
+	}
+	if back != v {
+		t.Errorf("back = %v; want %v", back, v)
+	}
+}