From a91c0519c23ddfe43c4b3ee8b3ec509da875070d Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Fri, 4 Sep 2026 10:29:49 +0000 Subject: [PATCH] change, mapper: distinguish deleted nodes Updates #3410 --- hscontrol/mapper/batcher.go | 15 +++-------- hscontrol/mapper/batcher_concurrency_test.go | 8 ++---- hscontrol/mapper/batcher_unit_test.go | 26 ++++++++++++++++++++ hscontrol/state/persist_test.go | 2 ++ hscontrol/types/change/change.go | 19 +++++++++++--- hscontrol/types/change/change_test.go | 23 +++++++++++++++++ 6 files changed, 73 insertions(+), 20 deletions(-) diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go index de3da87e0..9e6c2df8c 100644 --- a/hscontrol/mapper/batcher.go +++ b/hscontrol/mapper/batcher.go @@ -589,20 +589,13 @@ func (b *Batcher) addToBatch(changes ...change.Change) { // still has it registered. By cleaning up here, we prevent "node not found" // errors when workers try to generate map responses for deleted nodes. // - // Safety: [change.Change.PeersRemoved] is ONLY populated when nodes are actually - // deleted from the system (via [change.NodeRemoved] in [state.State.DeleteNode]). - // Policy changes that affect peer visibility do NOT use this field - they set - // RequiresRuntimePeerComputation=true and compute removed peers at runtime, - // putting them in [tailcfg.MapResponse.PeersRemoved] (a different struct). - // Therefore, this cleanup only removes nodes that are truly being deleted, - // not nodes that are still connected but have lost visibility of certain peers. - // This loop now also terminates the node's map sessions, so a future - // [change.Change.PeersRemoved] producer that is not a deletion would kill a - // live node's long poll, not merely evict its batcher entry. + // [change.Change.DeletedNodes] is an explicit internal lifecycle signal; + // [change.Change.PeersRemoved] remains only the protocol delta sent to clients + // when peers disappear from their view. // // See: https://github.com/juanfont/headscale/issues/2924 for _, ch := range changes { - for _, removedID := range ch.PeersRemoved { + for _, removedID := range ch.DeletedNodes { if nc, existed := b.nodes.LoadAndDelete(removedID); existed { b.totalNodes.Add(-1) diff --git a/hscontrol/mapper/batcher_concurrency_test.go b/hscontrol/mapper/batcher_concurrency_test.go index cbf53a811..3c21f122a 100644 --- a/hscontrol/mapper/batcher_concurrency_test.go +++ b/hscontrol/mapper/batcher_concurrency_test.go @@ -275,7 +275,7 @@ func TestAddToBatch_FullUpdateOverrides(t *testing.T) { }) } -// TestAddToBatch_NodeRemovalCleanup verifies that PeersRemoved in a change +// TestAddToBatch_NodeRemovalCleanup verifies that a permanent node deletion // cleans up the node from the batcher's internal state. func TestAddToBatch_NodeRemovalCleanup(t *testing.T) { lb := setupLightweightBatcher(t, 5, 10) @@ -287,11 +287,7 @@ func TestAddToBatch_NodeRemovalCleanup(t *testing.T) { _, exists := lb.b.nodes.Load(removedNode) require.True(t, exists, "node 3 should exist before removal") - // Send a change that includes node 3 in PeersRemoved - lb.b.addToBatch(change.Change{ - Reason: "node deleted", - PeersRemoved: []types.NodeID{removedNode}, - }) + lb.b.addToBatch(change.NodeRemoved(removedNode)) // Node should be removed from the nodes map _, exists = lb.b.nodes.Load(removedNode) diff --git a/hscontrol/mapper/batcher_unit_test.go b/hscontrol/mapper/batcher_unit_test.go index 378247261..5337ab179 100644 --- a/hscontrol/mapper/batcher_unit_test.go +++ b/hscontrol/mapper/batcher_unit_test.go @@ -455,6 +455,32 @@ func TestAddToBatch_NodeRemovedStopsSession(t *testing.T) { assert.Equal(t, int64(0), lb.b.totalNodes.Load()) } +func TestAddToBatch_PeersRemovedKeepsSession(t *testing.T) { + lb := setupLightweightBatcher(t, 1, 1) + defer lb.cleanup() + + mc, ok := lb.b.nodes.Load(1) + require.True(t, ok) + + stopped := make(chan struct{}) + + mc.mutex.Lock() + mc.connections[0].stop = func() { close(stopped) } + mc.mutex.Unlock() + + lb.b.AddWork(change.PeersRemoved(1)) + + select { + case <-stopped: + t.Fatal("a peer visibility delta must not stop the peer's own map session") + default: + } + + _, stillTracked := lb.b.nodes.Load(1) + assert.True(t, stillTracked, "a visible peer removal must remain tracked by the batcher") + assert.Equal(t, int64(1), lb.b.totalNodes.Load()) +} + // ============================================================================ // multiChannelNodeConn connection management Tests // ============================================================================ diff --git a/hscontrol/state/persist_test.go b/hscontrol/state/persist_test.go index 76623735b..def1ae736 100644 --- a/hscontrol/state/persist_test.go +++ b/hscontrol/state/persist_test.go @@ -532,6 +532,8 @@ func TestDeleteNodeReturnsRemovalOnPolicyFailure(t *testing.T) { require.ErrorIs(t, err, errInjectedPolicyNodeUpdate) assert.Equal(t, []types.NodeID{nodeID}, c.PeersRemoved, "a committed deletion must still notify peers and stop the node's session") + assert.Equal(t, []types.NodeID{nodeID}, c.DeletedNodes, + "a committed deletion must identify the session to stop") _, ok = s.GetNodeByID(nodeID) assert.False(t, ok, "a committed deletion must remove the in-memory node") diff --git a/hscontrol/types/change/change.go b/hscontrol/types/change/change.go index fc9401033..ac5e550b4 100644 --- a/hscontrol/types/change/change.go +++ b/hscontrol/types/change/change.go @@ -39,6 +39,11 @@ type Change struct { PeerPatches []*tailcfg.PeerChange SendAllPeers bool + // DeletedNodes identifies nodes permanently removed from state. Unlike + // PeersRemoved, this is an internal lifecycle signal: the batcher uses it + // to tear down the deleted nodes' own map sessions. + DeletedNodes []types.NodeID + // RequiresRuntimePeerComputation indicates that peer visibility // must be computed at runtime per-node. Used for policy changes // where each node may have different peer visibility. @@ -78,6 +83,7 @@ func (r Change) Merge(other Change) Change { merged.PeersChanged = uniqueNodeIDs(slices.Concat(r.PeersChanged, other.PeersChanged)) merged.PeersRemoved = uniqueNodeIDs(slices.Concat(r.PeersRemoved, other.PeersRemoved)) + merged.DeletedNodes = uniqueNodeIDs(slices.Concat(r.DeletedNodes, other.DeletedNodes)) merged.PeerPatches = slices.Concat(r.PeerPatches, other.PeerPatches) // Preserve [Change.OriginNode] for self-update detection. @@ -139,6 +145,7 @@ func (r Change) IsEmpty() bool { return len(r.PeersChanged) == 0 && len(r.PeersRemoved) == 0 && + len(r.DeletedNodes) == 0 && len(r.PeerPatches) == 0 } @@ -147,7 +154,8 @@ func (r Change) IsSelfOnly() bool { return false } - if r.SendAllPeers || len(r.PeersChanged) > 0 || len(r.PeersRemoved) > 0 || len(r.PeerPatches) > 0 { + if r.SendAllPeers || len(r.PeersChanged) > 0 || len(r.PeersRemoved) > 0 || + len(r.DeletedNodes) > 0 || len(r.PeerPatches) > 0 { return false } @@ -185,7 +193,8 @@ func (r Change) Type() string { return "patch" } - if len(r.PeersChanged) > 0 || len(r.PeersRemoved) > 0 || r.SendAllPeers { + if len(r.PeersChanged) > 0 || len(r.PeersRemoved) > 0 || + len(r.DeletedNodes) > 0 || r.SendAllPeers { return "peers" } @@ -438,7 +447,11 @@ func NodeAdded(id types.NodeID) Change { // NodeRemoved returns a [Change] for when a node is removed. func NodeRemoved(id types.NodeID) Change { - return PeersRemoved(id) + return Change{ + Reason: "node removed", + PeersRemoved: []types.NodeID{id}, + DeletedNodes: []types.NodeID{id}, + } } // KeyExpiryFor returns a [Change] for when a node's key expiry changes. diff --git a/hscontrol/types/change/change_test.go b/hscontrol/types/change/change_test.go index 44871ad0d..3f0d71a5c 100644 --- a/hscontrol/types/change/change_test.go +++ b/hscontrol/types/change/change_test.go @@ -88,6 +88,11 @@ func TestChange_IsEmpty(t *testing.T) { response: Change{PeersRemoved: []types.NodeID{1}}, want: false, }, + { + name: "DeletedNodes not empty", + response: Change{DeletedNodes: []types.NodeID{1}}, + want: false, + }, { name: "PeerPatches not empty", response: Change{PeerPatches: []*tailcfg.PeerChange{{}}}, @@ -149,6 +154,11 @@ func TestChange_IsSelfOnly(t *testing.T) { response: Change{TargetNode: 1, IncludeSelf: true, PeersRemoved: []types.NodeID{2}}, want: false, }, + { + name: "self only with DeletedNodes is not self only", + response: Change{TargetNode: 1, IncludeSelf: true, DeletedNodes: []types.NodeID{2}}, + want: false, + }, { name: "self only with PeerPatches is not self only", response: Change{TargetNode: 1, IncludeSelf: true, PeerPatches: []*tailcfg.PeerChange{{}}}, @@ -218,6 +228,12 @@ func TestChange_Merge(t *testing.T) { r2: Change{PeersRemoved: []types.NodeID{2, 3}}, want: Change{PeersRemoved: []types.NodeID{1, 2, 3}}, }, + { + name: "deleted nodes deduplicated", + r1: Change{DeletedNodes: []types.NodeID{1, 2}}, + r2: Change{DeletedNodes: []types.NodeID{2, 3}}, + want: Change{DeletedNodes: []types.NodeID{1, 2, 3}}, + }, { name: "peer patches concatenated", r1: Change{PeerPatches: []*tailcfg.PeerChange{{NodeID: 1}}}, @@ -471,6 +487,13 @@ func TestPeersRemoved(t *testing.T) { assert.Equal(t, []types.NodeID{1, 2, 3}, r.PeersRemoved) } +func TestNodeRemoved(t *testing.T) { + r := NodeRemoved(42) + assert.Equal(t, "node removed", r.Reason) + assert.Equal(t, []types.NodeID{42}, r.PeersRemoved) + assert.Equal(t, []types.NodeID{42}, r.DeletedNodes) +} + func TestPeerPatched(t *testing.T) { patch := &tailcfg.PeerChange{NodeID: 1} r := PeerPatched("endpoint change", patch)