From 2047ef35b81573f9b2fe341e258478dbb69b5c5d Mon Sep 17 00:00:00 2001 From: Kristoffer Dalby Date: Wed, 21 Oct 2026 11:09:11 +0000 Subject: [PATCH] hscontrol: wait for map session cleanup during shutdown Batcher shutdown cuts reconnect grace short to keep shutdown prompt. Wait for each session's Disconnect to persist LastSeen before the database closes. --- hscontrol/app.go | 9 ++- hscontrol/mapper/batcher.go | 5 ++ hscontrol/poll.go | 16 ++++-- hscontrol/poll_test.go | 112 ++++++++++++++++++++++++++++++++++-- 4 files changed, 130 insertions(+), 12 deletions(-) diff --git a/hscontrol/app.go b/hscontrol/app.go index c66639853..90580c8f1 100644 --- a/hscontrol/app.go +++ b/hscontrol/app.go @@ -105,6 +105,8 @@ type Headscale struct { authProvider AuthProvider mapBatcher *mapper.Batcher + // One per long-poll session, held through its Disconnect cleanup, so + // shutdown can close state after waiting on it. clientStreamsOpen sync.WaitGroup } @@ -1331,14 +1333,17 @@ func (h *Headscale) SetServerURLForTest(tb testing.TB, url string) { } // StartBatcherForTest initialises and starts the map response batcher. -// It registers a cleanup function on tb to stop the batcher. +// It registers a cleanup function on tb to stop the batcher and wait for map sessions. // It panics when called outside of tests. func (h *Headscale) StartBatcherForTest(tb testing.TB) { tb.Helper() h.mapBatcher = mapper.NewBatcherAndMapper(h.cfg, h.state) h.mapBatcher.Start() - tb.Cleanup(func() { h.mapBatcher.Close() }) + tb.Cleanup(func() { + h.mapBatcher.Close() + h.clientStreamsOpen.Wait() + }) } // MapBatcher returns the map response batcher (for test use). diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go index 325936836..f8efdc21d 100644 --- a/hscontrol/mapper/batcher.go +++ b/hscontrol/mapper/batcher.go @@ -478,6 +478,11 @@ func (b *Batcher) Start() { b.wg.Go(b.doWork) } +// Done is closed when batcher shutdown begins. +func (b *Batcher) Done() <-chan struct{} { + return b.done +} + func (b *Batcher) Close() { // Signal shutdown to all goroutines, only once. // Workers and queueWork both select on done, so closing it diff --git a/hscontrol/poll.go b/hscontrol/poll.go index d12ffa9c9..e4808e055 100644 --- a/hscontrol/poll.go +++ b/hscontrol/poll.go @@ -172,6 +172,10 @@ func (m *mapSession) serve() { func (m *mapSession) serveLongPoll() { m.log.Trace().Caller().Msg("long poll session started") + // Add/Done must precede cleanup: LIFO defers make shutdown's Wait cover Disconnect. + m.h.clientStreamsOpen.Add(1) + defer m.h.clientStreamsOpen.Done() + // connectGen is set by [state.State.Connect] below and captured by the deferred cleanup closure. // Each Connect acquires one live session in state; the cleanup must release // it with exactly one [state.State.Disconnect] call, in every exit path, or @@ -214,12 +218,18 @@ func (m *mapSession) serveLongPoll() { ticker := time.NewTicker(time.Second) defer ticker.Stop() + waitForReconnect: for range 10 { if m.h.mapBatcher.IsConnected(m.node.ID) { break } - <-ticker.C + select { + case <-m.h.mapBatcher.Done(): + // Reconnect grace has no purpose once the server is stopping. + break waitForReconnect + case <-ticker.C: + } } } @@ -246,10 +256,6 @@ func (m *mapSession) serveLongPoll() { m.log.Info().Caller().Str(zf.Chan, fmt.Sprintf("%p", m.ch)).Msg("node has disconnected") }() - // Set up the client stream - m.h.clientStreamsOpen.Add(1) - defer m.h.clientStreamsOpen.Done() - ctx, cancel := context.WithCancel(context.WithValue(m.ctx, nodeNameContextKey, m.node.Hostname)) defer cancel() diff --git a/hscontrol/poll_test.go b/hscontrol/poll_test.go index becd61f7c..4738d91c0 100644 --- a/hscontrol/poll_test.go +++ b/hscontrol/poll_test.go @@ -2,6 +2,7 @@ package hscontrol import ( "context" + "fmt" "io" "net/http" "net/http/httptest" @@ -528,17 +529,13 @@ func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t * }() t.Cleanup(func() { - dummyCh := make(chan *tailcfg.MapResponse, 1) - _ = app.mapBatcher.AddNode(node.ID, dummyCh, tailcfg.CapabilityVersion(100), nil) - cancel() + app.mapBatcher.Close() select { case <-serveDone: case <-time.After(2 * time.Second): } - - _ = app.mapBatcher.RemoveNode(node.ID, dummyCh) }) select { @@ -565,6 +562,12 @@ func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t * t.Fatal("expected the blocked write to eventually complete") } + require.Eventually(t, session.cancelChClosed.Load, 2*time.Second, 10*time.Millisecond, + "stale eviction must stop the session before batcher shutdown") + + // Shutdown skips reconnect grace; the stream wait now includes disconnect cleanup. + app.mapBatcher.Close() + assert.Eventually(t, func() bool { select { case <-streamsClosed: @@ -772,6 +775,105 @@ func TestDeletedNodeEndsStreamCleanly(t *testing.T) { } } +// TestLongPollShutdownWaitsForCleanup catches streams counted as closed before +// their disconnect cleanup finishes, including sessions already in reconnect grace. +func TestLongPollShutdownWaitsForCleanup(t *testing.T) { + t.Parallel() + + for _, inGrace := range []bool{false, true} { + t.Run(fmt.Sprintf("in_grace=%t", inGrace), func(t *testing.T) { + t.Parallel() + + app := createTestApp(t) + user := app.state.CreateUserForTest("poll-shutdown-user") + createdNode := app.state.CreateRegisteredNodeForTest(user, "poll-shutdown-node") + require.NoError(t, app.state.UpdatePolicyManagerUsersForTest()) + app.cfg.Tuning.NodeMapSessionBufferedChanSize = 1 + + app.mapBatcher.Close() + require.NoError(t, app.state.Close()) + + var err error + + app.state, err = state.NewState(app.cfg) + require.NoError(t, err) + + app.mapBatcher = mapper.NewBatcherAndMapper(app.cfg, app.state) + app.mapBatcher.Start() + + node, ok := app.state.GetNodeByID(createdNode.ID) + require.True(t, ok) + + ctx, cancel := context.WithCancel(t.Context()) + writer := newDelayedSuccessResponseWriter(0) + session := app.newMapSession(ctx, tailcfg.MapRequest{ + Stream: true, + Version: tailcfg.CapabilityVersion(100), + }, writer, node.AsStruct()) + + serveDone := make(chan struct{}) + + go func() { + session.serveLongPoll() + close(serveDone) + }() + + t.Cleanup(func() { + cancel() + app.mapBatcher.Close() + + select { + case <-serveDone: + case <-time.After(15 * time.Second): + t.Error("map session did not finish cleanup") + } + + require.NoError(t, app.state.Close()) + }) + + require.Eventually(t, func() bool { + return writer.WriteCount() >= 2 && app.mapBatcher.IsConnected(node.ID()) + }, 2*time.Second, 10*time.Millisecond, "map session should connect") + + if inGrace { + cancel() + require.Eventually(t, func() bool { + return !app.mapBatcher.IsConnected(node.ID()) + }, 2*time.Second, 10*time.Millisecond, "session should enter reconnect grace") + + pending, ok := app.state.GetNodeByID(node.ID()) + require.True(t, ok) + require.Equal(t, 1, pending.ActiveSessions(), "reconnect grace should keep the node online") + } + + streamsClosed := make(chan struct{}) + + go func() { + app.mapBatcher.Close() + app.clientStreamsOpen.Wait() + close(streamsClosed) + }() + + select { + case <-streamsClosed: + case <-time.After(3 * time.Second): + t.Fatal("shutdown must interrupt reconnect grace and wait for cleanup") + } + + disconnected, ok := app.state.GetNodeByID(node.ID()) + require.True(t, ok) + assert.Zero(t, disconnected.ActiveSessions(), "shutdown must wait for the session to release") + assert.False(t, disconnected.Online(), "the last session must take the node offline before shutdown returns") + + persisted, err := app.state.DB().GetNodeByID(node.ID()) + require.NoError(t, err) + require.NotNil(t, persisted.LastSeen) + assert.True(t, persisted.LastSeen.Equal(disconnected.LastSeen().Get()), + "shutdown must wait for the disconnect write to finish") + }) + } +} + func TestDeletedNodeInterruptsBlockedWrite(t *testing.T) { t.Parallel()