diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go index e21d86465..de3da87e0 100644 --- a/hscontrol/mapper/batcher.go +++ b/hscontrol/mapper/batcher.go @@ -596,12 +596,26 @@ func (b *Batcher) addToBatch(changes ...change.Change) { // 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. // // See: https://github.com/juanfont/headscale/issues/2924 for _, ch := range changes { for _, removedID := range ch.PeersRemoved { - if _, existed := b.nodes.LoadAndDelete(removedID); existed { + 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") diff --git a/hscontrol/mapper/batcher_unit_test.go b/hscontrol/mapper/batcher_unit_test.go index 08d515cdd..378247261 100644 --- a/hscontrol/mapper/batcher_unit_test.go +++ b/hscontrol/mapper/batcher_unit_test.go @@ -428,6 +428,33 @@ func TestMultiChannelClose_PreventsSendPanic(t *testing.T) { "send after close should return errConnectionClosed, not panic") } +func TestAddToBatch_NodeRemovedStopsSession(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.NodeRemoved(1)) + + // addToBatch runs synchronously, so the session must already be stopped. + select { + case <-stopped: + default: + t.Fatal("deleting a node must stop its map session, otherwise the long poll is orphaned") + } + + _, stillTracked := lb.b.nodes.Load(1) + assert.False(t, stillTracked, "deleted node should no longer be tracked by the batcher") + assert.Equal(t, int64(0), lb.b.totalNodes.Load()) +} + // ============================================================================ // multiChannelNodeConn connection management Tests // ============================================================================