state: reuse peer adjacency for payload-only writes

A write that cannot move visibility carries the previous adjacency
forward; policy and user changes rebuild it explicitly.

Updates #3417

(cherry picked from commit e3c4c81b18)

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
@@ -610,13 +663,17 @@ func resolveGivenName(nodes map[types.NodeID]types.Node, self types.NodeID, base
 	}
 }
 
-// snapshotFromNodes builds the index maps and primary-route table for
+// snapshotFromNodes builds a Snapshot from nodes. With reusePeers the
-// a new [Snapshot]. prevRoutes carries forward the previous primary
+// previous peer-ID adjacency is carried over unchanged; with reuseRoutes
-// assignment so a still-valid choice survives unrelated batches.
+// 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()
@@ -630,7 +687,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,
@@ -638,24 +722,12 @@ func snapshotFromNodes(
 		allNodes:          allNodes,
 		nodesByNodeKey:    make(map[key.NodePublic]types.NodeView),
 		nodesByMachineKey: make(map[key.MachinePublic]map[types.UserID]types.NodeView),
-
+		peersByNode:       peerIDsByNode,
-		// peersByNode is most likely the most expensive operation,
+		nodesByUser:       make(map[types.UserID][]types.NodeView),
-		// it will use the list of all nodes, combined with the
+		routes:            routes,
-		// current policy to precalculate which nodes are peers and
+		isPrimaryRoute:    isPrimaryRoute,
-		// 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,
 	}
 
-	// Build nodesByUser, nodesByNodeKey, and nodesByMachineKey maps
 	for _, n := range nodes {
 		nodeView := nodeViewsByID[n.ID]
 		userID := n.TypedUserID()
@@ -1022,7 +1094,11 @@ func (s *NodeStore) RebuildPeerMaps() {
 		rebuildResult: result,
 	}
 
-	s.writeQueue <- w
+	select {
+	case s.writeQueue <- w:
+	case <-s.stopped:
+		return
+	}
 
 	<-result
 }

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)
 		})
 	}
@@ -1383,6 +1385,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.
 //

hscontrol/state/state.go

@@ -1245,6 +1245,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
 }