net/tsaddr: use bart in NewContainsIPFunc, add tests, benchmarks

NewContainsIPFunc was previously documented as performing poorly if
there were many netip.Prefixes to search over. As such, we never it used it
in such cases.

This updates it to use bart at a certain threshold (over 6 prefixes,
currently), at which point the bart lookup overhead pays off.

This is currently kinda useless because we're not using it. But now we
can and get wins elsewhere. And we can remove the caveat in the docs.

    goos: darwin
    goarch: arm64
    pkg: tailscale.com/net/tsaddr
                                     │    before    │                after                 │
                                     │    sec/op    │    sec/op     vs base                │
    NewContainsIPFunc/empty-8          2.215n ± 11%   2.239n ±  1%   +1.08% (p=0.022 n=10)
    NewContainsIPFunc/cidr-list-1-8    17.44n ±  0%   17.59n ±  6%   +0.89% (p=0.000 n=10)
    NewContainsIPFunc/cidr-list-2-8    27.85n ±  0%   28.13n ±  1%   +1.01% (p=0.000 n=10)
    NewContainsIPFunc/cidr-list-3-8    36.05n ±  0%   36.56n ± 13%   +1.41% (p=0.000 n=10)
    NewContainsIPFunc/cidr-list-4-8    43.73n ±  0%   44.38n ±  1%   +1.50% (p=0.000 n=10)
    NewContainsIPFunc/cidr-list-5-8    51.61n ±  2%   51.75n ±  0%        ~ (p=0.101 n=10)
    NewContainsIPFunc/cidr-list-10-8   95.65n ±  0%   68.92n ±  0%  -27.94% (p=0.000 n=10)
    NewContainsIPFunc/one-ip-8         4.466n ±  0%   4.469n ±  1%        ~ (p=0.491 n=10)
    NewContainsIPFunc/two-ip-8         8.002n ±  1%   7.997n ±  4%        ~ (p=0.697 n=10)
    NewContainsIPFunc/three-ip-8       27.98n ±  1%   27.75n ±  0%   -0.82% (p=0.012 n=10)
    geomean                            19.60n         19.07n         -2.71%

Updates #12486

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

cmd/derper/depaware.txt

@@ -6,10 +6,12 @@ tailscale.com/cmd/derper dependencies: (generated by github.com/tailscale/depawa
    W    github.com/alexbrainman/sspi/internal/common                 from github.com/alexbrainman/sspi/negotiate
    W 💣 github.com/alexbrainman/sspi/negotiate                       from tailscale.com/net/tshttpproxy
         github.com/beorn7/perks/quantile                             from github.com/prometheus/client_golang/prometheus
+        github.com/bits-and-blooms/bitset                            from github.com/gaissmai/bart
      💣 github.com/cespare/xxhash/v2                                 from github.com/prometheus/client_golang/prometheus
    L    github.com/coreos/go-iptables/iptables                       from tailscale.com/util/linuxfw
    W 💣 github.com/dblohm7/wingoes                                   from tailscale.com/util/winutil
         github.com/fxamacker/cbor/v2                                 from tailscale.com/tka
+        github.com/gaissmai/bart                                     from tailscale.com/net/tsaddr
         github.com/golang/groupcache/lru                             from tailscale.com/net/dnscache
    L    github.com/google/nftables                                   from tailscale.com/util/linuxfw
    L 💣 github.com/google/nftables/alignedbuff                       from github.com/google/nftables/xt

cmd/stund/depaware.txt

@@ -1,7 +1,9 @@
 tailscale.com/cmd/stund dependencies: (generated by github.com/tailscale/depaware)
 
         github.com/beorn7/perks/quantile                             from github.com/prometheus/client_golang/prometheus
+        github.com/bits-and-blooms/bitset                            from github.com/gaissmai/bart
      💣 github.com/cespare/xxhash/v2                                 from github.com/prometheus/client_golang/prometheus
+        github.com/gaissmai/bart                                     from tailscale.com/net/tsaddr
         github.com/google/uuid                                       from tailscale.com/util/fastuuid
      💣 github.com/prometheus/client_golang/prometheus               from tailscale.com/tsweb/promvarz
         github.com/prometheus/client_golang/prometheus/internal      from github.com/prometheus/client_golang/prometheus

cmd/tailscale/depaware.txt

@@ -5,10 +5,12 @@ tailscale.com/cmd/tailscale dependencies: (generated by github.com/tailscale/dep
    W 💣 github.com/alexbrainman/sspi                                 from github.com/alexbrainman/sspi/internal/common+
    W    github.com/alexbrainman/sspi/internal/common                 from github.com/alexbrainman/sspi/negotiate
    W 💣 github.com/alexbrainman/sspi/negotiate                       from tailscale.com/net/tshttpproxy
+        github.com/bits-and-blooms/bitset                            from github.com/gaissmai/bart
    L    github.com/coreos/go-iptables/iptables                       from tailscale.com/util/linuxfw
    W 💣 github.com/dblohm7/wingoes                                   from github.com/dblohm7/wingoes/pe+
    W 💣 github.com/dblohm7/wingoes/pe                                from tailscale.com/util/winutil/authenticode
         github.com/fxamacker/cbor/v2                                 from tailscale.com/tka
+        github.com/gaissmai/bart                                     from tailscale.com/net/tsaddr
         github.com/golang/groupcache/lru                             from tailscale.com/net/dnscache
    L    github.com/google/nftables                                   from tailscale.com/util/linuxfw
    L 💣 github.com/google/nftables/alignedbuff                       from github.com/google/nftables/xt

net/tsaddr/tsaddr.go

@@ -11,6 +11,7 @@ import (
 	"slices"
 	"sync"
 
+	"github.com/gaissmai/bart"
 	"go4.org/netipx"
 	"tailscale.com/net/netaddr"
 	"tailscale.com/types/views"
@@ -165,6 +166,10 @@ func FalseContainsIPFunc() func(ip netip.Addr) bool {
 	return func(ip netip.Addr) bool { return false }
 }
 
+// pathForTest is a test hook for NewContainsIPFunc, to test that it took the
+// right construction path.
+var pathForTest = func(string) {}
+
 // NewContainsIPFunc returns a func that reports whether ip is in addrs.
 //
 // It's optimized for the cases of addrs being empty and addrs
@@ -176,32 +181,50 @@ func NewContainsIPFunc(addrs views.Slice[netip.Prefix]) func(ip netip.Addr) bool
 	// Specialize the three common cases: no address, just IPv4
 	// (or just IPv6), and both IPv4 and IPv6.
 	if addrs.Len() == 0 {
+		pathForTest("empty")
 		return func(netip.Addr) bool { return false }
 	}
-	// If any addr is more than a single IP, then just do the slow
+	// If any addr is a prefix with more than a single IP, then do either a
-	// linear thing until
+	// linear scan or a bart table, depending on the number of addrs.
-	// https://github.com/inetaf/netaddr/issues/139 is done.
 	if addrs.ContainsFunc(func(p netip.Prefix) bool { return !p.IsSingleIP() }) {
-		acopy := addrs.AsSlice()
+		if addrs.Len() > 6 {
-		return func(ip netip.Addr) bool {
+			pathForTest("bart")
-			for _, a := range acopy {
+			// Built a bart table.
-				if a.Contains(ip) {
+			t := &bart.Table[struct{}]{}
-					return true
+			for i := range addrs.Len() {
-				}
+				t.Insert(addrs.At(i), struct{}{})
+			}
+			return func(ip netip.Addr) bool {
+				_, ok := t.Get(ip)
+				return ok
+			}
+		} else {
+			pathForTest("linear-contains")
+			// Small enough to do a linear search.
+			acopy := addrs.AsSlice()
+			return func(ip netip.Addr) bool {
+				for _, a := range acopy {
+					if a.Contains(ip) {
+						return true
+					}
+				}
+				return false
 			}
-			return false
 		}
 	}
 	// Fast paths for 1 and 2 IPs:
 	if addrs.Len() == 1 {
+		pathForTest("one-ip")
 		a := addrs.At(0)
 		return func(ip netip.Addr) bool { return ip == a.Addr() }
 	}
 	if addrs.Len() == 2 {
+		pathForTest("two-ip")
 		a, b := addrs.At(0), addrs.At(1)
 		return func(ip netip.Addr) bool { return ip == a.Addr() || ip == b.Addr() }
 	}
 	// General case:
+	pathForTest("ip-map")
 	m := map[netip.Addr]bool{}
 	for i := range addrs.Len() {
 		m[addrs.At(i).Addr()] = true

net/tsaddr/tsaddr_test.go

@@ -8,6 +8,7 @@ import (
 	"testing"
 
 	"tailscale.com/net/netaddr"
+	"tailscale.com/tstest"
 	"tailscale.com/types/views"
 )
 
@@ -66,32 +67,140 @@ func TestCGNATRange(t *testing.T) {
 	}
 }
 
+func pp(ss ...string) (ret []netip.Prefix) {
+	for _, s := range ss {
+		ret = append(ret, netip.MustParsePrefix(s))
+	}
+	return
+}
+
+func aa(ss ...string) (ret []netip.Addr) {
+	for _, s := range ss {
+		ret = append(ret, netip.MustParseAddr(s))
+	}
+	return
+}
+
+var newContainsIPFuncTests = []struct {
+	name    string
+	pfx     []netip.Prefix
+	want    string
+	wantIn  []netip.Addr
+	wantOut []netip.Addr
+}{
+	{
+		name:    "empty",
+		pfx:     pp(),
+		want:    "empty",
+		wantOut: aa("8.8.8.8"),
+	},
+	{
+		name:    "cidr-list-1",
+		pfx:     pp("10.0.0.0/8"),
+		want:    "linear-contains",
+		wantIn:  aa("10.0.0.1", "10.2.3.4"),
+		wantOut: aa("8.8.8.8"),
+	},
+	{
+		name:    "cidr-list-2",
+		pfx:     pp("1.0.0.0/8", "3.0.0.0/8"),
+		want:    "linear-contains",
+		wantIn:  aa("1.0.0.1", "3.0.0.1"),
+		wantOut: aa("2.0.0.1"),
+	},
+	{
+		name:    "cidr-list-3",
+		pfx:     pp("1.0.0.0/8", "3.0.0.0/8", "5.0.0.0/8"),
+		want:    "linear-contains",
+		wantIn:  aa("1.0.0.1", "5.0.0.1"),
+		wantOut: aa("2.0.0.1"),
+	},
+	{
+		name:    "cidr-list-4",
+		pfx:     pp("1.0.0.0/8", "3.0.0.0/8", "5.0.0.0/8", "7.0.0.0/8"),
+		want:    "linear-contains",
+		wantIn:  aa("1.0.0.1", "7.0.0.1"),
+		wantOut: aa("2.0.0.1"),
+	},
+	{
+		name:    "cidr-list-5",
+		pfx:     pp("1.0.0.0/8", "3.0.0.0/8", "5.0.0.0/8", "7.0.0.0/8", "9.0.0.0/8"),
+		want:    "linear-contains",
+		wantIn:  aa("1.0.0.1", "9.0.0.1"),
+		wantOut: aa("2.0.0.1"),
+	},
+	{
+		name: "cidr-list-10",
+		pfx: pp("1.0.0.0/8", "3.0.0.0/8", "5.0.0.0/8", "7.0.0.0/8", "9.0.0.0/8",
+			"11.0.0.0/8", "13.0.0.0/8", "15.0.0.0/8", "17.0.0.0/8", "19.0.0.0/8"),
+		want:    "bart", // big enough that bart is faster than linear-contains
+		wantIn:  aa("1.0.0.1", "19.0.0.1"),
+		wantOut: aa("2.0.0.1"),
+	},
+	{
+		name:    "one-ip",
+		pfx:     pp("10.1.0.0/32"),
+		want:    "one-ip",
+		wantIn:  aa("10.1.0.0"),
+		wantOut: aa("10.0.0.9"),
+	},
+	{
+		name:    "two-ip",
+		pfx:     pp("10.1.0.0/32", "10.2.0.0/32"),
+		want:    "two-ip",
+		wantIn:  aa("10.1.0.0", "10.2.0.0"),
+		wantOut: aa("8.8.8.8"),
+	},
+	{
+		name:    "three-ip",
+		pfx:     pp("10.1.0.0/32", "10.2.0.0/32", "10.3.0.0/32"),
+		want:    "ip-map",
+		wantIn:  aa("10.1.0.0", "10.2.0.0"),
+		wantOut: aa("8.8.8.8"),
+	},
+}
+
+func BenchmarkNewContainsIPFunc(b *testing.B) {
+	for _, tt := range newContainsIPFuncTests {
+		b.Run(tt.name, func(b *testing.B) {
+			f := NewContainsIPFunc(views.SliceOf(tt.pfx))
+			for i := 0; i < b.N; i++ {
+				for _, ip := range tt.wantIn {
+					if !f(ip) {
+						b.Fatal("unexpected false")
+					}
+				}
+				for _, ip := range tt.wantOut {
+					if f(ip) {
+						b.Fatal("unexpected true")
+					}
+				}
+			}
+		})
+	}
+}
+
 func TestNewContainsIPFunc(t *testing.T) {
-	f := NewContainsIPFunc(views.SliceOf([]netip.Prefix{netip.MustParsePrefix("10.0.0.0/8")}))
+	for _, tt := range newContainsIPFuncTests {
-	if f(netip.MustParseAddr("8.8.8.8")) {
+		t.Run(tt.name, func(t *testing.T) {
-		t.Fatal("bad")
+			var got string
-	}
+			tstest.Replace(t, &pathForTest, func(path string) { got = path })
-	if !f(netip.MustParseAddr("10.1.2.3")) {
+
-		t.Fatal("bad")
+			f := NewContainsIPFunc(views.SliceOf(tt.pfx))
-	}
+			if got != tt.want {
-	f = NewContainsIPFunc(views.SliceOf([]netip.Prefix{netip.MustParsePrefix("10.1.2.3/32")}))
+				t.Errorf("func type = %q; want %q", got, tt.want)
-	if !f(netip.MustParseAddr("10.1.2.3")) {
+			}
-		t.Fatal("bad")
+			for _, ip := range tt.wantIn {
-	}
+				if !f(ip) {
-	f = NewContainsIPFunc(views.SliceOf([]netip.Prefix{
+					t.Errorf("match(%v) = false; want true", ip)
-		netip.MustParsePrefix("10.1.2.3/32"),
+				}
-		netip.MustParsePrefix("::2/128"),
+			}
-	}))
+			for _, ip := range tt.wantOut {
-	if !f(netip.MustParseAddr("::2")) {
+				if f(ip) {
-		t.Fatal("bad")
+					t.Errorf("match(%v) = true; want false", ip)
-	}
+				}
-	f = NewContainsIPFunc(views.SliceOf([]netip.Prefix{
+			}
-		netip.MustParsePrefix("10.1.2.3/32"),
+		})
-		netip.MustParsePrefix("10.1.2.4/32"),
-		netip.MustParsePrefix("::2/128"),
-	}))
-	if !f(netip.MustParseAddr("10.1.2.4")) {
-		t.Fatal("bad")
 	}
 }