From f1499ce280fa574b88380b444d035785bce497fe Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Sat, 24 Oct 2026 07:08:59 +0000 Subject: [PATCH] mapper: avoid refiltering broadcast recipients Backport 648905480 from #3568. --- hscontrol/mapper/batcher.go | 10 +- hscontrol/mapper/batcher_bench_test.go | 182 +++++++++++++++++++ hscontrol/mapper/batcher_concurrency_test.go | 121 ++++++++++++ 3 files changed, 306 insertions(+), 7 deletions(-) diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go index d61b9a893..e37daa067 100644 --- a/hscontrol/mapper/batcher.go +++ b/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 }) diff --git a/hscontrol/mapper/batcher_bench_test.go b/hscontrol/mapper/batcher_bench_test.go index b210a69f6..b95fd5e63 100644 --- a/hscontrol/mapper/batcher_bench_test.go +++ b/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") + }) + } + } + } + } +} diff --git a/hscontrol/mapper/batcher_concurrency_test.go b/hscontrol/mapper/batcher_concurrency_test.go index 59f3ef134..ed0d3c211 100644 --- a/hscontrol/mapper/batcher_concurrency_test.go +++ b/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) {