state,policy: add peer map and NodeStore write benchmarks

Cover BuildPeerMap, SetNodes, NodeStore writes and UpdateNodeFromMapRequest across node counts and policy shapes.
2026-10-06 06:40:06 +09:00 · 2026-09-25 17:36:26 +00:00
parent d813144f3b
commit 86b4f7c430
3 changed files with 379 additions and 12 deletions
@@ -1,6 +1,7 @@
 package v2

 import (
+	"fmt"
 	"net/netip"
 	"slices"
 	"testing"
@@ -2729,3 +2730,142 @@ func TestTagOwnedByTags(t *testing.T) {
 		require.False(t, nilPM.TagOwnedByTags("tag:leaf", []string{"tag:root"}))
 	})
 }
+
+// benchPolicies are the ACL shapes BenchmarkBuildPeerMap and
+// BenchmarkSetNodes measure: a global ACL, autogroup:self, and a via
+// grant. hscontrol/state/node_store_test.go checks the same three
+// shapes for adjacency correctness; this is its own copy since test
+// packages can't share one. Route dsts here are 10.0.0.0/8, wider than
+// state's fixed 10.33.0.0/24: benchNodes below spreads its ~5% of
+// routed nodes across 10.0.0.0/24 .. 10.255.0.0/24, so the dst must
+// cover that whole range for the route-owning ACL/grant rule to be
+// exercised rather than silently matching nothing.
+var benchPolicies = []struct {
+	name   string
+	policy string
+	// routerFiltered is whether the first routed node (node 1) gets a
+	// filter rule, so the route-owning rule is known to be exercised.
+	routerFiltered bool
+}{
+	{name: "global", policy: `{
+		"groups": {"group:a": ["u1@"]},
+		"tagOwners": {"tag:srv": ["u1@"]},
+		"acls": [
+			{"action": "accept", "src": ["group:a"], "dst": ["tag:srv:*"]},
+			{"action": "accept", "src": ["u2@"], "dst": ["10.0.0.0/8:*"]}
+		]}`, routerFiltered: true},
+	{name: "self", policy: `{
+		"acls": [{"action": "accept", "src": ["autogroup:member"], "dst": ["autogroup:self:*"]}]}`},
+	{name: "via", policy: `{
+		"tagOwners": {"tag:router": ["u1@"]},
+		"grants": [{"src": ["u2@"], "dst": ["10.0.0.0/8"], "ip": ["*"], "via": ["tag:router"]}]}`, routerFiltered: true},
+}
+
+// benchSetNodesPolicies is the subset of benchPolicies BenchmarkSetNodes
+// covers: via's per-node filter cost is close enough to global's that a
+// third axis wouldn't add signal.
+var benchSetNodesPolicies = benchPolicies[:2]
+
+// benchNodes builds n nodes spread across 10 users for the peer-map
+// benchmarks below: ~10% tagged tag:srv, ~5% also tagged tag:router and
+// carrying an approved and announced 10.x.0.0/24 route, each with a
+// unique IPv4.
+func benchNodes(n int) ([]types.User, types.Nodes) {
+	users := make([]types.User, 10)
+	for i := range users {
+		users[i] = types.User{ID: uint(i + 1), Name: fmt.Sprintf("u%d", i+1)}
+	}
+
+	nodes := make(types.Nodes, 0, n)
+	for i := range n {
+		u := users[i%len(users)]
+		ip := netip.AddrFrom4([4]byte{100, 64, byte(i / 256), byte(i % 256)}) //nolint:gosec
+
+		nd := &types.Node{
+			ID:       types.NodeID(i + 1),
+			Hostname: fmt.Sprintf("n%d", i+1),
+			IPv4:     &ip,
+		}
+
+		if i%10 == 0 {
+			nd.Tags = []string{"tag:srv"}
+		} else {
+			nd.UserID, nd.User = &u.ID, &u
+		}
+
+		if i%20 == 0 {
+			// The via policy's routers.
+			nd.Tags = []string{"tag:router", "tag:srv"}
+			subnet := netip.PrefixFrom(netip.AddrFrom4([4]byte{10, byte((i / 20) % 256), 0, 0}), 24) //nolint:gosec
+			nd.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: []netip.Prefix{subnet}}
+			nd.ApprovedRoutes = []netip.Prefix{subnet}
+		}
+
+		nodes = append(nodes, nd)
+	}
+
+	return users, nodes
+}
+
+// BenchmarkBuildPeerMap measures PolicyManager.BuildPeerMap over
+// realistic node counts and policy shapes, without any NodeStore or
+// reuse-path involvement: the number this baseline compares later
+// NodeStore write-path changes against.
+func BenchmarkBuildPeerMap(b *testing.B) {
+	for _, pol := range benchPolicies {
+		for _, n := range []int{100, 300, 617, 1000} {
+			b.Run(fmt.Sprintf("%s/n=%d", pol.name, n), func(b *testing.B) {
+				users, nodes := benchNodes(n)
+				pm, err := NewPolicyManager([]byte(pol.policy), users, nodes.ViewSlice())
+				require.NoError(b, err)
+
+				rules, err := pm.FilterForNode(nodes[0].View())
+				require.NoError(b, err)
+				require.Equal(b, pol.routerFiltered, len(rules) > 0,
+					"router filter rules: %v", rules)
+
+				b.ReportAllocs()
+
+				for b.Loop() {
+					pm.BuildPeerMap(nodes.ViewSlice())
+				}
+			})
+		}
+	}
+}
+
+// BenchmarkSetNodes measures PolicyManager.SetNodes when one node's
+// route changes on every call, the same per-write cost
+// BenchmarkNodeStoreWrite drives through a NodeStore.
+func BenchmarkSetNodes(b *testing.B) {
+	for _, pol := range benchSetNodesPolicies {
+		b.Run(fmt.Sprintf("%s/n=%d", pol.name, 617), func(b *testing.B) {
+			users, nodes := benchNodes(617)
+			pm, err := NewPolicyManager([]byte(pol.policy), users, nodes.ViewSlice())
+			require.NoError(b, err)
+
+			b.ReportAllocs()
+
+			i := 0
+			for b.Loop() {
+				i++
+				subnet := netip.PrefixFrom(netip.AddrFrom4([4]byte{10, byte(200 + i%50), 0, 0}), 24) //nolint:gosec
+				// The policy manager holds views of the previous node; mutating
+				// it in place would change both sides of SetNodes' comparison.
+				nd := nodes[0].Clone()
+				nd.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: []netip.Prefix{subnet}}
+				nd.ApprovedRoutes = []netip.Prefix{subnet}
+				nodes[0] = nd
+
+				changed, err := pm.SetNodes(nodes.ViewSlice())
+				if err != nil {
+					b.Fatal(err)
+				}
+
+				if !changed {
+					b.Fatal("SetNodes must see the route change and recompile")
+				}
+			}
+		})
+	}
+}
@@ -1,6 +1,7 @@
 package state

 import (
+	"fmt"
 	"net/netip"
 	"strings"
 	"sync"
@@ -8,6 +9,7 @@ import (
 	"testing"
 	"time"

+	"github.com/juanfont/headscale/hscontrol/db"
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/types/change"
 	"github.com/stretchr/testify/assert"
@@ -1008,3 +1010,110 @@ func TestBuildMapRequestChangeResponse(t *testing.T) {
 		})
 	}
 }
+
+// benchMapRequestSetup pre-creates a sqlite database with n registered
+// nodes spread across 10 users (~10% tagged tag:srv, half of those also
+// carrying an approved and announced 10.x.0.0/24 route), then constructs a State
+// that loads them, at benchmark scale the same way persistTestSetup
+// does for a single node. Returns the State and the ID of node 0, a
+// plain node with neither tag nor route, to drive requests against.
+func benchMapRequestSetup(b *testing.B, n int) (*State, types.NodeID) {
+	b.Helper()
+
+	dbPath := b.TempDir() + "/headscale.db"
+	cfg := persistTestConfig(dbPath)
+
+	database, err := db.NewHeadscaleDatabase(cfg)
+	require.NoError(b, err)
+
+	users := make([]*types.User, 10)
+	for i := range users {
+		users[i] = database.CreateUserForTest(fmt.Sprintf("u%d", i+1))
+	}
+
+	var targetID types.NodeID
+
+	for i := range n {
+		node := database.CreateRegisteredNodeForTest(users[i%len(users)], fmt.Sprintf("n%d", i))
+
+		if i == 0 {
+			targetID = node.ID
+
+			continue
+		}
+
+		if i%10 != 0 {
+			continue
+		}
+
+		node.Tags = []string{"tag:srv"}
+
+		if i%20 == 0 {
+			subnet := netip.PrefixFrom(netip.AddrFrom4([4]byte{10, byte((i / 20) % 256), 0, 0}), 24) //nolint:gosec
+			node.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: []netip.Prefix{subnet}}
+			node.ApprovedRoutes = []netip.Prefix{subnet}
+		}
+
+		require.NoError(b, database.DB.Save(node).Error)
+	}
+
+	require.NoError(b, database.Close())
+
+	s, err := NewState(cfg)
+	require.NoError(b, err)
+	b.Cleanup(func() { _ = s.Close() })
+
+	target, ok := s.GetNodeByID(targetID)
+	require.True(b, ok)
+	require.False(b, target.IsTagged(), "target must be plain")
+	require.Empty(b, target.AnnouncedRoutes(), "target must be plain")
+
+	routers := 0
+
+	for _, nv := range s.ListNodes().All() {
+		if len(nv.AnnouncedRoutes()) > 0 {
+			routers++
+		}
+	}
+
+	require.Positive(b, routers, "setup must create subnet routers")
+
+	return s, targetID
+}
+
+// BenchmarkUpdateNodeFromMapRequest measures the no-op path
+// TestNoOpMapRequestSkipsPersist and TestNoOpMapRequestEmitsNoPeerChange
+// pin the behaviour of: a MapRequest that is value-identical to the
+// node's current state, against a realistic node count. The baseline
+// later NodeStore write-path changes are compared against.
+func BenchmarkUpdateNodeFromMapRequest(b *testing.B) {
+	b.Run("identical/n=617", func(b *testing.B) {
+		s, nodeID := benchMapRequestSetup(b, 617)
+
+		nv, ok := s.GetNodeByID(nodeID)
+		require.True(b, ok, "target node should exist in NodeStore")
+
+		req := tailcfg.MapRequest{
+			NodeKey:  nv.NodeKey(),
+			DiscoKey: nv.DiscoKey(),
+			Hostinfo: &tailcfg.Hostinfo{
+				Hostname: nv.Hostname(),
+				NetInfo:  &tailcfg.NetInfo{PreferredDERP: 1},
+			},
+		}
+
+		// Establish the Hostinfo/DERP state once so every request timed
+		// below is a genuine no-op.
+		_, err := s.UpdateNodeFromMapRequest(nodeID, req)
+		require.NoError(b, err)
+
+		b.ReportAllocs()
+
+		for b.Loop() {
+			_, err := s.UpdateNodeFromMapRequest(nodeID, req)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
@@ -1735,6 +1735,27 @@ func checkAdjacencyMatchesFullBuild(t fatalfer, store *NodeStore, pm *policyv2.P
 	}
 }

+// Policy shapes shared by TestNodeStoreAdjacencyMatchesFullBuild and the
+// NodeStore write benchmarks below: a global ACL, autogroup:self, and a
+// via grant. hscontrol/policy/v2/policy_test.go keeps its own copy since
+// test packages can't share one.
+const (
+	policyGlobal = `{
+		"groups": {"group:a": ["u1@"]},
+		"tagOwners": {"tag:srv": ["u1@"]},
+		"acls": [
+			{"action": "accept", "src": ["group:a"], "dst": ["tag:srv:*"]},
+			{"action": "accept", "src": ["u2@"], "dst": ["10.33.0.0/24:*"]}
+		]}`
+
+	policyAutogroupSelf = `{
+		"acls": [{"action": "accept", "src": ["autogroup:member"], "dst": ["autogroup:self:*"]}]}`
+
+	policyVia = `{
+		"tagOwners": {"tag:router": ["u1@"]},
+		"grants": [{"src": ["u2@"], "dst": ["10.33.0.0/24"], "ip": ["*"], "via": ["tag:router"]}]}`
+)
+
 // TestNodeStoreAdjacencyMatchesFullBuild drives random node mutations
 // through a real [NodeStore] wired to a real [policyv2.PolicyManager] and
 // checks, after every step, that the resulting adjacency — including
@@ -1753,18 +1774,9 @@ func TestNodeStoreAdjacencyMatchesFullBuild(t *testing.T) {
 		name string
 		pol  string
 	}{
-		{name: "global", pol: `{
-			"groups": {"group:a": ["u1@"]},
-			"tagOwners": {"tag:srv": ["u1@"]},
-			"acls": [
-				{"action": "accept", "src": ["group:a"], "dst": ["tag:srv:*"]},
-				{"action": "accept", "src": ["u2@"], "dst": ["10.33.0.0/24:*"]}
-			]}`},
-		{name: "autogroup-self", pol: `{
-			"acls": [{"action": "accept", "src": ["autogroup:member"], "dst": ["autogroup:self:*"]}]}`},
-		{name: "via", pol: `{
-			"tagOwners": {"tag:router": ["u1@"]},
-			"grants": [{"src": ["u2@"], "dst": ["10.33.0.0/24"], "ip": ["*"], "via": ["tag:router"]}]}`},
+		{name: "global", pol: policyGlobal},
+		{name: "autogroup-self", pol: policyAutogroupSelf},
+		{name: "via", pol: policyVia},
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			rapid.Check(t, func(rt *rapid.T) {
@@ -1919,3 +1931,109 @@ func TestUpdateChangesReportsRelationFields(t *testing.T) {
 		})
 	}
 }
+
+// benchNodes builds n nodes spread across 10 users for
+// BenchmarkNodeStoreWrite: ~10% tagged tag:srv, ~5% carrying an
+// approved and announced 10.x.0.0/24 route, each with a unique IPv4.
+// hscontrol/policy/v2/policy_test.go keeps its own copy since test
+// packages can't share one.
+func benchNodes(n int) ([]types.User, types.Nodes) {
+	users := make([]types.User, 10)
+	for i := range users {
+		users[i] = types.User{ID: uint(i + 1), Name: fmt.Sprintf("u%d", i+1)}
+	}
+
+	nodes := make(types.Nodes, 0, n)
+	for i := range n {
+		u := users[i%len(users)]
+		nd := createTestNode(types.NodeID(i+1), u.ID, u.Name, fmt.Sprintf("n%d", i+1))
+
+		ip := netip.AddrFrom4([4]byte{100, 64, byte(i / 256), byte(i % 256)}) //nolint:gosec
+		nd.IPv4, nd.IPv6 = &ip, nil
+
+		if i%10 == 0 {
+			nd.Tags = []string{"tag:srv"}
+		} else {
+			nd.User = &u
+		}
+
+		if i%20 == 0 {
+			subnet := netip.PrefixFrom(netip.AddrFrom4([4]byte{10, byte((i / 20) % 256), 0, 0}), 24) //nolint:gosec
+			nd.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: []netip.Prefix{subnet}}
+			nd.ApprovedRoutes = []netip.Prefix{subnet}
+		}
+
+		nodes = append(nodes, &nd)
+	}
+
+	return users, nodes
+}
+
+// BenchmarkNodeStoreWrite drives one UpdateNode of the named kind
+// followed by syncPolicy against a started NodeStore wired to a real
+// PolicyManager, over a realistic node count. peer-builds/op counts
+// calls into the wrapped peersFunc, the baseline later NodeStore
+// write-path changes are compared against: a payload-only write
+// (lastseen) should cost far fewer builds than a relation-changing one
+// (route, tag).
+func BenchmarkNodeStoreWrite(b *testing.B) {
+	users, nodes := benchNodes(617)
+
+	for _, kind := range []struct {
+		name   string
+		mutate func(i int, n *types.Node)
+	}{
+		{name: "lastseen", mutate: func(_ int, n *types.Node) {
+			n.LastSeen = new(time.Now())
+		}},
+		{name: "route", mutate: func(i int, n *types.Node) {
+			subnet := netip.PrefixFrom(netip.AddrFrom4([4]byte{10, byte(200 + i%50), 0, 0}), 24) //nolint:gosec
+			n.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: []netip.Prefix{subnet}}
+			n.ApprovedRoutes = []netip.Prefix{subnet}
+		}},
+		{name: "tag", mutate: func(i int, n *types.Node) {
+			tag := "tag:srv"
+			if i%2 == 0 {
+				tag = "tag:web"
+			}
+
+			n.Tags = []string{tag}
+			n.UserID, n.User = nil, nil
+		}},
+	} {
+		b.Run(fmt.Sprintf("%s/n=%d", kind.name, len(nodes)), func(b *testing.B) {
+			var calls atomic.Int64
+
+			pm, err := policyv2.NewPolicyManager([]byte(policyGlobal), users, nodes.ViewSlice())
+			require.NoError(b, err)
+
+			peersFunc := func(ns []types.NodeView) map[types.NodeID][]types.NodeID {
+				calls.Add(1)
+
+				return pm.BuildPeerMap(views.SliceOf(ns))
+			}
+
+			store := NewNodeStore(nodes, peersFunc, TestBatchSize, TestBatchTimeout)
+
+			store.Start()
+			defer store.Stop()
+
+			targetID := nodes[0].ID
+
+			// NewNodeStore's own initial build is setup, not a per-write cost.
+			calls.Store(0)
+
+			b.ReportAllocs()
+
+			i := 0
+			for b.Loop() {
+				i++
+
+				store.UpdateNode(targetID, func(n *types.Node) { kind.mutate(i, n) })
+				syncPolicy(b, store, pm)
+			}
+
+			b.ReportMetric(float64(calls.Load())/float64(b.N), "peer-builds/op")
+		})
+	}
+}