diff --git a/hscontrol/state/node_store.go b/hscontrol/state/node_store.go index 050c979e..08d4ad0a 100644 --- a/hscontrol/state/node_store.go +++ b/hscontrol/state/node_store.go @@ -89,6 +89,13 @@ var ( Name: "nodestore_queue_depth", Help: "Current depth of NodeStore write queue", }) + + // Bounded labels only: no node IDs or free-form reasons. + nodeStoreSnapshotBuilds = promauto.NewCounterVec(prometheus.CounterOpts{ + Namespace: prometheusNamespace, + Name: "nodestore_snapshot_builds_total", + Help: "Snapshot builds by how peer adjacency was obtained: recomputed through the policy or reused from the previous snapshot.", + }, []string{"peers"}) ) // NodeStore is a thread-safe store for nodes. @@ -124,7 +131,7 @@ func NewNodeStore(allNodes types.Nodes, peersFunc PeersFunc, batchSize int, batc nodes[n.ID] = *n } - snap := snapshotFromNodes(nodes, peersFunc, nil) + snap := snapshotFromNodes(nodes, peersFunc, nil, false, false) store := &NodeStore{ peersFunc: peersFunc, @@ -193,6 +200,22 @@ type work struct { multiUpdates map[types.NodeID]UpdateNodeFunc } +// updateChanges reports whether an in-place update moved a peer-visibility +// input (which also re-elects routes) or only a route-election input. +// Peer visibility depends on what the policy reads; election additionally +// depends on online and health state, treating unknown online as offline. +func updateChanges(pre, post *types.Node) (bool, bool) { + preView, postView := pre.View(), post.View() + if postView.HasPolicyChange(preView) || postView.HasNetworkChanges(preView) { + return true, true + } + + wasOnline := pre.IsOnline != nil && *pre.IsOnline + isOnline := post.IsOnline != nil && *post.IsOnline + + return false, wasOnline != isOnline || pre.Unhealthy != post.Unhealthy +} + // PutNode adds or updates a node in the store. // If the node already exists, it will be replaced. // If the node does not exist, it will be added. @@ -450,6 +473,12 @@ func (s *NodeStore) applyBatch(batch []work) { // NodeView for that work. setErrResults := make(map[*work]error) + // relationChanged forces a peersFunc run; electionChanged forces a + // route re-election. put/del/setName/rebuildPeerMaps are treated as + // relation-changing; updateMulti compares the node before and after. + relationChanged := false + electionChanged := false + for i := range batch { w := &batch[i] switch w.op { @@ -461,6 +490,9 @@ func (s *NodeStore) applyBatch(batch []work) { if w.nodeResult != nil { nodeResultRequests[w.nodeID] = append(nodeResultRequests[w.nodeID], w) } + + relationChanged = true + electionChanged = true case updateMulti: for id, fn := range w.multiUpdates { n, exists := nodes[id] @@ -469,6 +501,7 @@ func (s *NodeStore) applyBatch(batch []work) { } oldGivenName := n.GivenName + pre := n.Clone() fn(&n) if n.GivenName != oldGivenName { @@ -476,6 +509,10 @@ func (s *NodeStore) applyBatch(batch []work) { } nodes[id] = n + + relation, election := updateChanges(pre, &n) + relationChanged = relationChanged || relation + electionChanged = electionChanged || election } case del: delete(nodes, w.nodeID) @@ -483,6 +520,9 @@ func (s *NodeStore) applyBatch(batch []work) { if w.nodeResult != nil { nodeResultRequests[w.nodeID] = append(nodeResultRequests[w.nodeID], w) } + + relationChanged = true + electionChanged = true case setName: n, exists := nodes[w.nodeID] if !exists { @@ -518,15 +558,28 @@ func (s *NodeStore) applyBatch(batch []work) { n.GivenName = w.name nodes[w.nodeID] = n nodeResultRequests[w.nodeID] = append(nodeResultRequests[w.nodeID], w) + relationChanged = true + electionChanged = true case rebuildPeerMaps: // rebuildPeerMaps doesn't modify nodes, it just forces the snapshot rebuild // below to recalculate peer relationships using the current peersFunc rebuildOps = append(rebuildOps, w) + relationChanged = true + electionChanged = true } } prev := s.data.Load() - newSnap := snapshotFromNodes(nodes, s.peersFunc, prev.routes) + + // A relation change recomputes adjacency; anything else reuses the + // previous peer IDs and re-elects routes only if an election input moved. + if relationChanged { + nodeStoreSnapshotBuilds.WithLabelValues("recomputed").Inc() + } else { + nodeStoreSnapshotBuilds.WithLabelValues("reused").Inc() + } + + newSnap := snapshotFromNodes(nodes, s.peersFunc, prev, !relationChanged, !electionChanged) s.data.Store(&newSnap) // Update node count gauge @@ -601,13 +654,17 @@ func resolveGivenName(nodes map[types.NodeID]types.Node, self types.NodeID, base } } -// snapshotFromNodes builds the index maps and primary-route table for -// a new [Snapshot]. prevRoutes carries forward the previous primary -// assignment so a still-valid choice survives unrelated batches. +// snapshotFromNodes builds a Snapshot from nodes. With reusePeers the +// previous peer-ID adjacency is carried over unchanged; with reuseRoutes +// the previous route election is. ListPeers resolves adjacency IDs +// through this snapshot's fresh views, so a reused adjacency never +// serves stale node payloads. prev may be nil only when both reuse +// flags are false. func snapshotFromNodes( nodes map[types.NodeID]types.Node, peersFunc PeersFunc, - prevRoutes map[netip.Prefix]types.NodeID, + prev *Snapshot, + reusePeers, reuseRoutes bool, ) Snapshot { timer := prometheus.NewTimer(nodeStoreSnapshotBuildDuration) defer timer.ObserveDuration() @@ -621,7 +678,34 @@ func snapshotFromNodes( nodeViewsByID[n.ID] = nv } - routes, isPrimaryRoute := electPrimaryRoutes(nodes, prevRoutes) + var ( + routes map[netip.Prefix]types.NodeID + isPrimaryRoute map[types.NodeID]bool + ) + + if reuseRoutes { + routes, isPrimaryRoute = prev.routes, prev.isPrimaryRoute + } else { + // Carrying the previous assignment forward lets a still-valid + // primary survive unrelated batches. + var prevRoutes map[netip.Prefix]types.NodeID + if prev != nil { + prevRoutes = prev.routes + } + + routes, isPrimaryRoute = electPrimaryRoutes(nodes, prevRoutes) + } + + var peerIDsByNode map[types.NodeID][]types.NodeID + + if reusePeers { + peerIDsByNode = prev.peersByNode + } else { + peersTimer := prometheus.NewTimer(nodeStorePeersCalculationDuration) + peerIDsByNode = peersFunc(allNodes) + + peersTimer.ObserveDuration() + } newSnap := Snapshot{ nodesByID: nodes, @@ -629,24 +713,12 @@ func snapshotFromNodes( allNodes: allNodes, nodesByNodeKey: make(map[key.NodePublic]types.NodeView), nodesByMachineKey: make(map[key.MachinePublic]map[types.UserID]types.NodeView), - - // peersByNode is most likely the most expensive operation, - // it will use the list of all nodes, combined with the - // current policy to precalculate which nodes are peers and - // can see each other. - peersByNode: func() map[types.NodeID][]types.NodeID { - peersTimer := prometheus.NewTimer(nodeStorePeersCalculationDuration) - defer peersTimer.ObserveDuration() - - return peersFunc(allNodes) - }(), - nodesByUser: make(map[types.UserID][]types.NodeView), - - routes: routes, - isPrimaryRoute: isPrimaryRoute, + peersByNode: peerIDsByNode, + nodesByUser: make(map[types.UserID][]types.NodeView), + routes: routes, + isPrimaryRoute: isPrimaryRoute, } - // Build nodesByUser, nodesByNodeKey, and nodesByMachineKey maps for _, n := range nodes { nodeView := nodeViewsByID[n.ID] userID := n.TypedUserID() @@ -1006,7 +1078,11 @@ func (s *NodeStore) RebuildPeerMaps() { rebuildResult: result, } - s.writeQueue <- w + select { + case s.writeQueue <- w: + case <-s.stopped: + return + } <-result } diff --git a/hscontrol/state/node_store_test.go b/hscontrol/state/node_store_test.go index 210e27c1..1887c138 100644 --- a/hscontrol/state/node_store_test.go +++ b/hscontrol/state/node_store_test.go @@ -6,6 +6,7 @@ import ( "net/netip" "runtime" "sync" + "sync/atomic" "testing" "time" @@ -13,6 +14,7 @@ import ( "github.com/juanfont/headscale/hscontrol/types" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "tailscale.com/tailcfg" "tailscale.com/types/key" ) @@ -151,7 +153,7 @@ func TestSnapshotFromNodes(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { nodes, peersFunc := tt.setupFunc() - snapshot := snapshotFromNodes(nodes, peersFunc, nil) + snapshot := snapshotFromNodes(nodes, peersFunc, nil, false, false) tt.validate(t, nodes, snapshot) }) } @@ -1363,6 +1365,266 @@ func TestGetNodesByMachineKeyAllUsers(t *testing.T) { }) } +// TestPeerIrrelevantWriteReusesPeerMap ensures writes that cannot alter peer +// visibility neither run peersFunc nor copy the immutable adjacency map. +// +// peersByNode is derived from addresses, ownership, routes, tags, and exit-node +// status. LastSeen and node keys are payload/index data, so neither can change +// adjacency. +func TestPeerIrrelevantWriteReusesPeerMap(t *testing.T) { + var peersCalls atomic.Int64 + + countingPeersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeID { + peersCalls.Add(1) + + return allowAllPeersFunc(nodes) + } + + node1 := createTestNode(1, 1, "user1", "node1") + node2 := createTestNode(2, 2, "user2", "node2") + + store := NewNodeStore(types.Nodes{&node1, &node2}, countingPeersFunc, TestBatchSize, TestBatchTimeout) + store.Start() + + defer store.Stop() + + // Ignore the initial snapshot build. + peersCalls.Store(0) + + before := store.data.Load() + require.NotEmpty(t, before.peersByNode[1]) + + now := time.Now() + _, ok := store.UpdateNode(1, func(n *types.Node) { + n.LastSeen = &now + }) + require.True(t, ok, "update should apply") + + newNodeKey := key.NewNode().Public() + _, ok = store.UpdateNode(1, func(n *types.Node) { + n.NodeKey = newNodeKey + }) + require.True(t, ok, "key rotation should apply") + + indexed, ok := store.GetNodeByNodeKey(newNodeKey) + require.True(t, ok, "rotated key must be present in the rebuilt key index") + require.Equal(t, types.NodeID(1), indexed.ID()) + + peersOf2 := store.ListPeers(2) + require.Equal(t, 1, peersOf2.Len()) + require.Equal(t, newNodeKey, peersOf2.At(0).NodeKey(), + "reused adjacency must resolve to the fresh view") + + require.Equalf(t, int64(0), peersCalls.Load(), + "payload/index-only writes must not recompute the peer map, got %d recomputations", + peersCalls.Load()) + + _, ok = store.UpdateNode(1, func(n *types.Node) { + n.User = nil + }) + require.True(t, ok, "user association update should apply") + require.Equal(t, int64(1), peersCalls.Load(), + "a BuildPeerMap input must recompute peer adjacency") +} + +// TestHealthOnlyWriteReusesPeerMap ensures a health flip re-elects routes +// without recomputing peer adjacency. +func TestHealthOnlyWriteReusesPeerMap(t *testing.T) { + var peersCalls atomic.Int64 + + countingPeersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeID { + peersCalls.Add(1) + + return allowAllPeersFunc(nodes) + } + + // Set up two HA candidates for the same prefix. + node1 := createTestNode(1, 1, "user1", "router1") + node2 := createTestNode(2, 1, "user1", "router2") + + pfx := netip.MustParsePrefix("10.99.0.0/24") + node1.Hostinfo = &tailcfg.Hostinfo{Hostname: "router1", RoutableIPs: []netip.Prefix{pfx}} + node2.Hostinfo = &tailcfg.Hostinfo{Hostname: "router2", RoutableIPs: []netip.Prefix{pfx}} + node1.ApprovedRoutes = append(node1.ApprovedRoutes, pfx) + node2.ApprovedRoutes = append(node2.ApprovedRoutes, pfx) + + online := true + node1.IsOnline = &online + node2.IsOnline = &online + + store := NewNodeStore(types.Nodes{&node1, &node2}, countingPeersFunc, TestBatchSize, TestBatchTimeout) + store.Start() + + defer store.Stop() + + primary, ok := store.PrimaryRouteFor(pfx) + require.True(t, ok) + require.Equal(t, types.NodeID(1), primary) + + peersCalls.Store(0) // ignore initial snapshot build + + // Healthy -> healthy (no-op): no election, no relation rebuild. + _, ok = store.UpdateNode(1, func(n *types.Node) { + // Simulate BatchSetNodeHealth setter semantics with the same + // stored value. healthSetter(healthy=true) sets Unhealthy=false; + // node already has Unhealthy=false. + healthSetter(true)(n) + }) + require.True(t, ok) + + // Healthy -> unhealthy (real transition): election must run, but + // relation must NOT be recomputed (Unhealthy is election-relevant, + // not relation-relevant). + _, ok = store.UpdateNode(1, healthSetter(false)) + require.True(t, ok) + primary, ok = store.PrimaryRouteFor(pfx) + require.True(t, ok) + require.Equal(t, types.NodeID(2), primary) + + // Unhealthy -> unhealthy (no-op): no relation rebuild. + _, ok = store.UpdateNode(1, healthSetter(false)) + require.True(t, ok) + + require.Equal(t, int64(0), peersCalls.Load(), + "no health-only write may recompute the peer map; got %d recomputations", + peersCalls.Load()) +} + +func BenchmarkSnapshotPayloadDense(b *testing.B) { + const nodeCount = 500 + + nodes := make(map[types.NodeID]types.Node, nodeCount) + for i := 1; i <= nodeCount; i++ { + id := types.NodeID(i) //nolint:gosec // bounded benchmark node count + nodes[id] = createTestNode(id, uint(i), "user", "node") //nolint:gosec // bounded benchmark node count + } + + initial := snapshotFromNodes(nodes, allowAllPeersFunc, nil, false, false) + n := nodes[1] + n.LastSeen = new(time.Now()) + nodes[1] = n + + b.Run("reuse-peer-adjacency", func(b *testing.B) { + previous := initial + + b.ReportAllocs() + b.ResetTimer() + + for b.Loop() { + next := snapshotFromNodes(nodes, allowAllPeersFunc, &previous, true, true) + previous = next + } + }) + + b.Run("rebuild-peer-adjacency", func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + + for b.Loop() { + snapshotFromNodes(nodes, allowAllPeersFunc, nil, false, false) + } + }) +} + +// TestRebuildPeerMapsAfterStopReturns ensures a rebuild requested after the +// writer has exited does not block the caller forever. +func TestRebuildPeerMapsAfterStopReturns(t *testing.T) { + node := createTestNode(1, 1, "user1", "node1") + store := NewNodeStore(types.Nodes{&node}, allowAllPeersFunc, TestBatchSize, TestBatchTimeout) + store.Start() + store.Stop() + + done := make(chan struct{}) + + go func() { + store.RebuildPeerMaps() + close(done) + }() + + require.Eventually(t, func() bool { + select { + case <-done: + return true + default: + return false + } + }, 5*time.Second, 10*time.Millisecond, "RebuildPeerMaps hung after Stop") +} + +// TestUpdateNodeRecomputesPeersOnlyForRelationInputs pins which fields make a +// write recompute peer adjacency: the inputs that force a peer-map rebuild +// (an announced but unapproved route is included on purpose). +func TestUpdateNodeRecomputesPeersOnlyForRelationInputs(t *testing.T) { + subnet := netip.MustParsePrefix("10.77.0.0/24") + + tests := []struct { + name string + mutate func(*types.Node) + wantRecompute bool + }{ + {name: "last seen", mutate: func(n *types.Node) { n.LastSeen = new(time.Now()) }}, + {name: "node key", mutate: func(n *types.Node) { n.NodeKey = key.NewNode().Public() }}, + {name: "expiry", mutate: func(n *types.Node) { n.Expiry = new(time.Now()) }}, + {name: "online", mutate: func(n *types.Node) { n.IsOnline = new(true) }}, + {name: "unhealthy", mutate: func(n *types.Node) { n.Unhealthy = true }}, + { + name: "endpoints", + mutate: func(n *types.Node) { + n.Endpoints = []netip.AddrPort{netip.MustParseAddrPort("203.0.113.1:41641")} + }, + }, + {name: "tags", mutate: func(n *types.Node) { n.Tags = []string{"tag:x"} }, wantRecompute: true}, + { + name: "ipv4", + mutate: func(n *types.Node) { + ip := netip.MustParseAddr("100.64.9.9") + n.IPv4 = &ip + }, + wantRecompute: true, + }, + { + name: "announced route", + mutate: func(n *types.Node) { + n.Hostinfo = &tailcfg.Hostinfo{RoutableIPs: []netip.Prefix{subnet}} + }, + wantRecompute: true, + }, + {name: "user association", mutate: func(n *types.Node) { n.User = nil }, wantRecompute: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var peersCalls atomic.Int64 + + countingPeersFunc := func(nodes []types.NodeView) map[types.NodeID][]types.NodeID { + peersCalls.Add(1) + + return allowAllPeersFunc(nodes) + } + + node1 := createTestNode(1, 1, "user1", "node1") + node2 := createTestNode(2, 2, "user2", "node2") + + store := NewNodeStore(types.Nodes{&node1, &node2}, countingPeersFunc, TestBatchSize, TestBatchTimeout) + store.Start() + + defer store.Stop() + + peersCalls.Store(0) + + _, ok := store.UpdateNode(1, tt.mutate) + require.True(t, ok) + + var want int64 + if tt.wantRecompute { + want = 1 + } + + require.Equal(t, want, peersCalls.Load()) + }) + } +} + // TestListPeersExcludesSelf proves a node is never returned among its own // peers, on both the snapshot path and the explicit peer-ID path. // diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go index 1a321707..2c29ffa5 100644 --- a/hscontrol/state/state.go +++ b/hscontrol/state/state.go @@ -1203,6 +1203,12 @@ func (s *State) SetPolicy(pol []byte) (bool, error) { // Clear SSH check auth times when policy changes. s.ClearSSHCheckAuth() + // Payload-only writes reuse the cached adjacency, so a policy swap + // must rebuild it here rather than wait for the next relation write. + if changed { + s.nodeStore.RebuildPeerMaps() + } + return changed, nil }