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

@@ -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).

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

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()
 

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()