From 552356a2ee81a4dc34898a4c1ffe1c35467415ba Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Thu, 22 Oct 2026 01:56:00 +0000
Subject: [PATCH] state: skip the caller's policy refresh once the writer
 applied it

The writer refreshes the policy manager before publishing its batch, so
the caller's SetNodes with the published snapshot recompiled it back.
---
 hscontrol/state/node_store_test.go | 42 ++++++++++++++++++++++++++++++
 hscontrol/state/persist_test.go    | 10 +++++++
 hscontrol/state/state.go           |  9 +++++++
 3 files changed, 61 insertions(+)

diff --git a/hscontrol/state/node_store_test.go b/hscontrol/state/node_store_test.go
index 6bdf7dd94..9c41b5ac2 100644
--- a/hscontrol/state/node_store_test.go
+++ b/hscontrol/state/node_store_test.go
@@ -2214,6 +2214,48 @@ func TestStatePolicyWriteBuildsPeersOnce(t *testing.T) {
 	}
 }
 
+// TestStateRefreshDoesNotMovePolicyBack replays a writer that has already
+// moved the policy manager to the batch it is building, a snapshot not yet
+// published, when a caller whose write landed earlier refreshes the policy
+// from the published one. The caller must not move the policy manager back:
+// that recompiles, and the peer map rebuild it forces recompiles forward
+// again, so under churn every write paid for three compiles.
+func TestStateRefreshDoesNotMovePolicyBack(t *testing.T) {
+	s, ids, builds := newPeerBuildTestState(t)
+
+	genBefore := s.polMan.NodesGeneration()
+
+	_, ok := s.nodeStore.UpdateNode(ids[1], func(n *types.Node) { n.Tags = []string{"tag:a"} })
+	require.True(t, ok)
+
+	// The next batch, built but not yet published.
+	ahead := make(types.Nodes, 0, s.nodeStore.ListNodes().Len())
+	for _, nv := range s.nodeStore.ListNodes().All() {
+		ahead = append(ahead, nv.AsStruct())
+	}
+
+	for _, n := range ahead {
+		if n.ID == ids[2] {
+			n.Tags = []string{"tag:b"}
+		}
+	}
+
+	changed, err := s.polMan.SetNodes(ahead.ViewSlice())
+	require.NoError(t, err)
+	require.True(t, changed, "test sanity: the batch in flight moves the policy")
+
+	genAhead := s.polMan.NodesGeneration()
+
+	builds.Store(0)
+
+	c, err := s.updatePolicyManagerNodes(genBefore)
+	require.NoError(t, err)
+
+	assert.Equal(t, "policy", c.Type(), "the caller's write still moved the policy")
+	assert.Equal(t, genAhead, s.polMan.NodesGeneration(), "the caller's refresh must not recompile")
+	assert.Equal(t, int64(0), builds.Load(), "the caller's refresh must not force a peer map rebuild")
+}
+
 // TestStateConcurrentTagWritesEachReportPolicyChange runs two SetNodeTags
 // on different nodes at once. The NodeStore may apply both in one batch,
 // so the policy manager sees both tags in a single SetNodes; each caller
diff --git a/hscontrol/state/persist_test.go b/hscontrol/state/persist_test.go
index ad781a00b..d8d1d9c8f 100644
--- a/hscontrol/state/persist_test.go
+++ b/hscontrol/state/persist_test.go
@@ -606,7 +606,17 @@ func TestDeleteNodeReturnsRemovalOnPolicyFailure(t *testing.T) {
 	node, ok := s.GetNodeByID(nodeID)
 	require.True(t, ok)
 
+	// Fail the writer's refresh too: it shares the policy manager, and a
+	// refresh it applied is not retried by the caller.
+	nodes := make(types.Nodes, 0, s.nodeStore.ListNodes().Len())
+	for _, nv := range s.nodeStore.ListNodes().All() {
+		nodes = append(nodes, nv.AsStruct())
+	}
+
+	s.nodeStore.Stop()
 	s.polMan = failingSetNodesPolicyManager{PolicyManager: s.polMan}
+	s.nodeStore = NewNodeStore(nodes, policyPeersFunc(s.polMan), TestBatchSize, TestBatchTimeout)
+	s.nodeStore.Start()
 
 	cs, err := s.DeleteNode(node)
 	require.ErrorIs(t, err, errInjectedPolicyNodeUpdate)
diff --git a/hscontrol/state/state.go b/hscontrol/state/state.go
index 7536b778b..152ecf643 100644
--- a/hscontrol/state/state.go
+++ b/hscontrol/state/state.go
@@ -3212,6 +3212,15 @@ func (s *State) DrainSelfRefreshes() []change.Change {
 // have the list already available so it could go much quicker. Alternatively
 // the policy manager could have a remove or add list for nodes.
 func (s *State) updatePolicyManagerNodes(genBefore uint64) (change.Change, error) {
+	// A move since genBefore is the writer applying this write's nodes or a
+	// later batch's, which it hands the policy manager before publishing.
+	// The published snapshot read here can then be older, and SetNodes with
+	// it would recompile the policy back. What is left to do here is retry a
+	// recompile the writer failed, which leaves the generation where it was.
+	if s.polMan.NodesGeneration() != genBefore {
+		return change.PolicyChange(), nil
+	}
+
 	nodes := s.ListNodes()
 
 	changed, err := s.polMan.SetNodes(nodes)