From 86b4f7c43052d86f6dc4baf0273263755ba0082e Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Fri, 25 Sep 2026 17:36:26 +0000 Subject: [PATCH] state,policy: add peer map and NodeStore write benchmarks Cover BuildPeerMap, SetNodes, NodeStore writes and UpdateNodeFromMapRequest across node counts and policy shapes. --- hscontrol/policy/v2/policy_test.go | 140 ++++++++++++++++++++++++++++ hscontrol/state/maprequest_test.go | 109 ++++++++++++++++++++++ hscontrol/state/node_store_test.go | 142 ++++++++++++++++++++++++++--- 3 files changed, 379 insertions(+), 12 deletions(-) diff --git a/hscontrol/policy/v2/policy_test.go b/hscontrol/policy/v2/policy_test.go index 9d6567e61..77c49a2b0 100644 --- a/hscontrol/policy/v2/policy_test.go +++ b/hscontrol/policy/v2/policy_test.go @@ -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") + } + } + }) + } +} diff --git a/hscontrol/state/maprequest_test.go b/hscontrol/state/maprequest_test.go index e437448e4..e448d48ff 100644 --- a/hscontrol/state/maprequest_test.go +++ b/hscontrol/state/maprequest_test.go @@ -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) + } + } + }) +} diff --git a/hscontrol/state/node_store_test.go b/hscontrol/state/node_store_test.go index 0c86e3526..053ab88ea 100644 --- a/hscontrol/state/node_store_test.go +++ b/hscontrol/state/node_store_test.go @@ -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") + }) + } +}