mapper: avoid refiltering broadcast recipients

Backport 648905480 from #3568.

hscontrol/mapper/batcher.go

@@ -734,18 +734,14 @@ func (b *Batcher) addToBatch(changes ...change.Change) {
 		}
 	}
 
-	// Handle broadcast changes - send to all nodes, filtering as needed
+	// Broadcasts have no TargetNode, so every node receives the same changes.
 	if len(broadcast) > 0 {
-		b.nodes.Range(func(nodeID types.NodeID, nc *multiChannelNodeConn) bool {
+		b.nodes.Range(func(_ types.NodeID, nc *multiChannelNodeConn) bool {
 			if nc == nil {
 				return true
 			}
 
-			filtered := change.FilterForNode(nodeID, broadcast)
-
-			if len(filtered) > 0 {
-				nc.appendPending(filtered...)
-			}
+			nc.appendPending(broadcast...)
 
 			return true
 		})

hscontrol/mapper/batcher_bench_test.go

@@ -17,7 +17,10 @@ import (
 
 	"github.com/juanfont/headscale/hscontrol/types"
 	"github.com/juanfont/headscale/hscontrol/types/change"
+	"github.com/juanfont/headscale/hscontrol/util/zlog/zf"
 	"github.com/puzpuzpuz/xsync/v4"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
 	"tailscale.com/tailcfg"
 )
 
@@ -740,3 +743,182 @@ func BenchmarkMapResponseFromChange(b *testing.B) {
 		})
 	}
 }
+
+// Frozen from 1096abfc8; keep the body unchanged for differential tests and benchmarks.
+func baselineAddToBatch(b *Batcher, changes ...change.Change) {
+	// Clean up any nodes being permanently removed from the system.
+	//
+	// This handles the case where a node is deleted from state but the batcher
+	// still has it registered. By cleaning up here, we prevent "node not found"
+	// errors when workers try to generate map responses for deleted nodes.
+	//
+	// [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.DeletedNodes {
+			if nc, existed := b.nodes.LoadAndDelete(removedID); existed {
+				b.totalNodes.Add(-1)
+
+				// Tear the node's map sessions down. Dropping the entry alone
+				// leaves [mapSession.serveLongPoll] streaming to a node that no
+				// longer exists: [Batcher.Close] ranges b.nodes and can no
+				// longer reach it, so shutdown blocks on clientStreamsOpen, and
+				// the client keeps polling a node it should be re-authenticating.
+				// See: https://github.com/juanfont/headscale/issues/3410
+				if nc != nil {
+					nc.close()
+				}
+
+				log.Debug().
+					Uint64(zf.NodeID, removedID.Uint64()).
+					Msg("removed deleted node from batcher")
+			}
+		}
+	}
+
+	// A full update supersedes every state change pending or in this call,
+	// but not the pings addressed to a node; those follow the full.
+	if change.HasFull(changes) {
+		b.nodes.Range(func(_ types.NodeID, nc *multiChannelNodeConn) bool {
+			if nc == nil {
+				return true
+			}
+
+			nc.collapsePendingToFull(changes)
+
+			return true
+		})
+
+		return
+	}
+
+	broadcast, targeted := change.SplitTargetedAndBroadcast(changes)
+
+	// Handle targeted changes - send only to the specific node
+	for _, ch := range targeted {
+		if nc, ok := b.nodes.Load(ch.TargetNode); ok && nc != nil {
+			nc.appendPending(ch)
+		}
+	}
+
+	// Handle broadcast changes - send to all nodes, filtering as needed
+	if len(broadcast) > 0 {
+		b.nodes.Range(func(nodeID types.NodeID, nc *multiChannelNodeConn) bool {
+			if nc == nil {
+				return true
+			}
+
+			filtered := change.FilterForNode(nodeID, broadcast)
+
+			if len(filtered) > 0 {
+				nc.appendPending(filtered...)
+			}
+
+			return true
+		})
+	}
+}
+
+var batchIngressImplementations = []struct {
+	name string
+	add  func(*Batcher, ...change.Change)
+}{
+	{"baseline", baselineAddToBatch},
+	{"candidate", (*Batcher).addToBatch},
+}
+
+// BenchmarkBatchIngressPaired measures enqueue work only; fixture construction,
+// delivery-count checks, and draining pending are outside the timer.
+func BenchmarkBatchIngressPaired(b *testing.B) {
+	level := zerolog.GlobalLevel()
+
+	zerolog.SetGlobalLevel(zerolog.Disabled)
+	defer zerolog.SetGlobalLevel(level)
+
+	for _, nodes := range []int{100, 1000, 10000} {
+		for _, batch := range []int{1, 8, 64} {
+			for _, mode := range []struct {
+				name  string
+				mixed bool
+				burst int
+			}{
+				{"broadcast", false, 1},
+				{"mixed", true, 1},
+				{"burst", false, 4},
+			} {
+				for _, impl := range batchIngressImplementations {
+					name := fmt.Sprintf(
+						"nodes=%d/batch=%d/%s/%s",
+						nodes,
+						batch,
+						mode.name,
+						impl.name,
+					)
+					b.Run(name, func(b *testing.B) {
+						b.StopTimer()
+
+						batcher, _ := benchBatcher(nodes, 1)
+						b.Cleanup(batcher.Close)
+
+						conns := make([]*multiChannelNodeConn, nodes)
+						for i := range conns {
+							conns[i], _ = batcher.nodes.Load(
+								types.NodeID(i + 1),
+							) //nolint:gosec // bounded benchmark
+						}
+
+						changes := make([]change.Change, 0, batch+2)
+						for i := range batch {
+							changes = append(
+								changes,
+								change.NodeOnline(types.NodeID(i+1)),
+							) //nolint:gosec // bounded benchmark
+						}
+
+						if mode.mixed {
+							changes = append(
+								changes,
+								change.SelfUpdate(1),
+								change.SelfUpdate(10001),
+							) //nolint:gosec // bounded benchmark
+						}
+
+						b.ReportAllocs()
+						b.ResetTimer()
+
+						for range b.N {
+							b.StartTimer()
+
+							for range mode.burst {
+								impl.add(batcher, changes...)
+							}
+
+							b.StopTimer()
+
+							for i, nc := range conns {
+								want := batch * mode.burst
+								if mode.mixed && i == 0 {
+									want += mode.burst
+								}
+
+								if got := len(nc.drainPending()); got != want {
+									b.Fatalf(
+										"node %d: got %d changes, want %d",
+										i+1,
+										got,
+										want,
+									)
+								}
+							}
+						}
+
+						b.ReportMetric(float64(mode.burst), "calls/op")
+					})
+				}
+			}
+		}
+	}
+}

hscontrol/mapper/batcher_concurrency_test.go

@@ -248,6 +248,127 @@ func TestAddToBatch_ConcurrentBroadcast(t *testing.T) {
 		"at least some nodes should have pending changes after broadcast")
 }
 
+func TestAddToBatchDifferential(t *testing.T) {
+	baseline := setupLightweightBatcher(t, 4, 1)
+	candidate := setupLightweightBatcher(t, 4, 1)
+	t.Cleanup(baseline.cleanup)
+	t.Cleanup(candidate.cleanup)
+
+	var stopped [2]atomic.Int64
+
+	for i, lb := range []*lightweightBatcher{baseline, candidate} {
+		nc, _ := lb.b.nodes.Load(4)
+		nc.connections[0].stop = func() { stopped[i].Add(1) }
+
+		lb.b.RemoveNode(3, lb.channels[3])
+		lb.b.nodes.Store(99, nil)
+	}
+
+	endpoint := change.EndpointOrDERPUpdate(1, &tailcfg.PeerChange{
+		NodeID: 1, DERPRegion: 7, Online: new(true),
+	})
+	ping := change.PingNode(2, &tailcfg.PingRequest{URL: "https://example.com/one"})
+	stages := []struct {
+		name    string
+		changes []change.Change
+	}{
+		{"empty", nil},
+		{"mixed", []change.Change{
+			change.NodeOnline(2), change.SelfUpdate(1), endpoint,
+			change.SelfUpdate(88), ping, change.NodeOnline(2),
+		}},
+		{"repeated_kinds", []change.Change{
+			change.NodeOnline(2), change.NodeOffline(2), change.NodeOnline(2),
+			change.PeersChanged("several peers", 3, 1),
+		}},
+		{"deletion", []change.Change{
+			change.NodeRemoved(4), change.SelfUpdate(4), change.PeersRemoved(3),
+		}},
+		{
+			"full_with_pings",
+			[]change.Change{
+				change.FullUpdate(),
+				change.PingNode(
+					1,
+					&tailcfg.PingRequest{URL: "https://example.com/two"},
+				),
+				change.PingNode(
+					88,
+					&tailcfg.PingRequest{URL: "https://example.com/missing"},
+				),
+				change.NodeOnline(3),
+			},
+		},
+		{"repeated_full", []change.Change{change.FullUpdate(), change.FullUpdate()}},
+		{"after_full", []change.Change{endpoint, change.SelfUpdate(1), ping}},
+	}
+
+	for _, stage := range stages {
+		t.Run(stage.name, func(t *testing.T) {
+			baselineAddToBatch(baseline.b, stage.changes...)
+			candidate.b.addToBatch(stage.changes...)
+
+			for _, id := range []types.NodeID{1, 2, 3, 4, 88, 99} {
+				want, wantExists := baseline.b.nodes.Load(id)
+				got, gotExists := candidate.b.nodes.Load(id)
+				require.Equal(t, wantExists, gotExists, "node %d", id)
+				require.Equal(t, want == nil, got == nil, "node %d", id)
+
+				if got == nil {
+					continue
+				}
+
+				assert.Equal(t, getPendingForNode(baseline.b, id),
+					getPendingForNode(candidate.b, id), "node %d", id)
+			}
+
+			assert.Equal(t, baseline.b.totalNodes.Load(), candidate.b.totalNodes.Load())
+			assert.Equal(t, stopped[0].Load(), stopped[1].Load())
+
+			if stage.name == "mixed" {
+				assert.Equal(t, []change.Change{
+					change.SelfUpdate(
+						1,
+					),
+					change.NodeOnline(2),
+					endpoint,
+					change.NodeOnline(2),
+				}, getPendingForNode(candidate.b, 1))
+				assert.Equal(t, []change.Change{
+					change.NodeOnline(2), endpoint, change.NodeOnline(2),
+				}, getPendingForNode(candidate.b, 3), "disconnected recipient retains broadcasts")
+			}
+		})
+	}
+
+	require.EqualValues(t, 1, stopped[1].Load(), "deletion must stop the session once")
+}
+
+func TestAddToBatchRecipientStorage(t *testing.T) {
+	for _, impl := range batchIngressImplementations {
+		t.Run(impl.name, func(t *testing.T) {
+			lb := setupLightweightBatcher(t, 2, 1)
+			t.Cleanup(lb.cleanup)
+
+			incoming := []change.Change{change.NodeOnline(1), change.NodeOffline(1)}
+			impl.add(lb.b, incoming...)
+			incoming[0] = change.FullUpdate()
+
+			first, _ := lb.b.nodes.Load(1)
+			first.pendingMu.Lock()
+			first.pending[0].Reason = "recipient-local change"
+			first.pendingMu.Unlock()
+
+			assert.Equal(
+				t,
+				[]change.Change{change.NodeOnline(1), change.NodeOffline(1)},
+				getPendingForNode(lb.b, 2),
+				"recipients must own their pending slices",
+			)
+		})
+	}
+}
+
 // TestAddToBatch_FullUpdateOverrides verifies that a FullUpdate replaces
 // all pending changes for every node.
 func TestAddToBatch_FullUpdateOverrides(t *testing.T) {