From 9c764e5a99ff43eac83f41a81c7382b541fb7fb1 Mon Sep 17 00:00:00 2001
From: Kristoffer Dalby <kristoffer@tailscale.com>
Date: Thu, 8 Oct 2026 09:54:04 +0000
Subject: [PATCH] state: return the replaced node from NodeStore updates

UpdateNodeDiff hands back the pre-update clone the writer already makes.

Updates #3531
---
 hscontrol/state/auth_tagged_expiry_test.go |  2 +-
 hscontrol/state/node_store.go              | 38 ++++++++++++++++++++++
 hscontrol/state/node_store_test.go         | 22 +++++++++++++
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/hscontrol/state/auth_tagged_expiry_test.go b/hscontrol/state/auth_tagged_expiry_test.go
index bd94a5d78..4398e3085 100644
--- a/hscontrol/state/auth_tagged_expiry_test.go
+++ b/hscontrol/state/auth_tagged_expiry_test.go
@@ -1822,7 +1822,7 @@ func waitParkedOnWriteQueue(t *testing.T, fn string) {
 		stacks := string(buf[:runtime.Stack(buf, true)])
 		for g := range strings.SplitSeq(stacks, "\n\n") {
 			if strings.Contains(g, "[select") &&
-				strings.Contains(g, "(*NodeStore).UpdateNodes(") &&
+				strings.Contains(g, "(*NodeStore).updateNodes(") &&
 				strings.Contains(g, "."+fn+"(") {
 				return true
 			}
diff --git a/hscontrol/state/node_store.go b/hscontrol/state/node_store.go
index 0f3e4c021..67672d2c8 100644
--- a/hscontrol/state/node_store.go
+++ b/hscontrol/state/node_store.go
@@ -198,6 +198,9 @@ type work struct {
 	// prober applying multiple probe results at once) cannot have a
 	// partial snapshot published between the updates.
 	multiUpdates map[types.NodeID]UpdateNodeFunc
+	// pre, when non-nil, receives each updated node as it was before its
+	// update function ran.
+	pre map[types.NodeID]types.NodeView
 }
 
 // updateChanges reports whether an in-place update moved a peer-visibility
@@ -274,6 +277,27 @@ func (s *NodeStore) UpdateNode(nodeID types.NodeID, updateFn UpdateNodeFunc) (ty
 	return s.GetNode(nodeID)
 }
 
+// UpdateNodeDiff is [NodeStore.UpdateNode] that also returns, first, the node
+// as it was before updateFn ran, so callers can tell what the update changed
+// without cloning the node themselves. That view is invalid when the node does
+// not exist.
+func (s *NodeStore) UpdateNodeDiff(
+	nodeID types.NodeID,
+	updateFn UpdateNodeFunc,
+) (types.NodeView, types.NodeView, bool) {
+	timer := prometheus.NewTimer(nodeStoreOperationDuration.WithLabelValues("update"))
+	defer timer.ObserveDuration()
+
+	pre := make(map[types.NodeID]types.NodeView, 1)
+	s.updateNodes(map[types.NodeID]UpdateNodeFunc{nodeID: updateFn}, pre)
+
+	nodeStoreOperations.WithLabelValues("update").Inc()
+
+	after, ok := s.GetNode(nodeID)
+
+	return pre[nodeID], after, ok
+}
+
 // UpdateNodes applies per-node update functions in a single atomic
 // batch. The election that recomputes primary routes runs once, after
 // every update has landed, so callers cannot observe an intermediate
@@ -282,6 +306,15 @@ func (s *NodeStore) UpdateNode(nodeID types.NodeID, updateFn UpdateNodeFunc) (ty
 // would change the election outcome — e.g. the HA prober applying
 // concurrent probe-timeout results.
 func (s *NodeStore) UpdateNodes(updates map[types.NodeID]UpdateNodeFunc) {
+	s.updateNodes(updates, nil)
+}
+
+// updateNodes queues updates as one batch entry and waits for it, filling pre
+// when it is non-nil.
+func (s *NodeStore) updateNodes(
+	updates map[types.NodeID]UpdateNodeFunc,
+	pre map[types.NodeID]types.NodeView,
+) {
 	timer := prometheus.NewTimer(nodeStoreOperationDuration.WithLabelValues("update_multi"))
 	defer timer.ObserveDuration()
 
@@ -292,6 +325,7 @@ func (s *NodeStore) UpdateNodes(updates map[types.NodeID]UpdateNodeFunc) {
 	w := work{
 		op:           updateMulti,
 		multiUpdates: updates,
+		pre:          pre,
 		result:       make(chan struct{}),
 	}
 
@@ -507,6 +541,10 @@ func (s *NodeStore) applyBatch(batch []work) {
 
 				nodes[id] = n
 
+				if w.pre != nil {
+					w.pre[id] = pre.View()
+				}
+
 				relation, election := updateChanges(pre, &n)
 				relationChanged = relationChanged || relation
 				electionChanged = electionChanged || election
diff --git a/hscontrol/state/node_store_test.go b/hscontrol/state/node_store_test.go
index 8c8a27501..6bdf7dd94 100644
--- a/hscontrol/state/node_store_test.go
+++ b/hscontrol/state/node_store_test.go
@@ -1653,6 +1653,28 @@ func TestUpdateNodeRecomputesPeersOnlyForRelationInputs(t *testing.T) {
 	}
 }
 
+// TestUpdateNodeDiff proves the update returns the node the writer replaced
+// alongside the result, and an invalid before for a missing node.
+func TestUpdateNodeDiff(t *testing.T) {
+	node := createTestNode(1, 1, "user1", "node1")
+	oldKey := node.NodeKey
+
+	store := NewNodeStore(types.Nodes{&node}, allowAllPeersFunc, TestBatchSize, TestBatchTimeout)
+	store.Start()
+
+	defer store.Stop()
+
+	newKey := key.NewNode().Public()
+	before, after, ok := store.UpdateNodeDiff(1, func(n *types.Node) { n.NodeKey = newKey })
+	require.True(t, ok)
+	assert.Equal(t, oldKey, before.NodeKey())
+	assert.Equal(t, newKey, after.NodeKey())
+
+	before, _, ok = store.UpdateNodeDiff(99, func(*types.Node) {})
+	assert.False(t, ok)
+	assert.False(t, before.Valid())
+}
+
 // TestListPeersExcludesSelf proves a node is never returned among its own
 // peers, on both the snapshot path and the explicit peer-ID path.
 //