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)
2026-10-26 00:06:53 +09:00 · 2026-09-23 21:19:22 +02:00
1 parent 938c2bd753
commit 237dc74e73
@@ -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.
 //