state, poll: refcount poll sessions, mark offline only on last release

A cancelled map request whose handler ran late could Connect after the live session, steal the newest SessionEpoch, then exit without disconnecting (stillConnected path); the live session's final Disconnect was rejected as stale and the node stayed online forever (relogin flake). Counted releases are order-independent, so overlapping sessions cannot strand a node in either direction.
2026-07-08 00:50:20 +09:00 · 2026-06-10 12:21:08 +00:00
parent 759381ad78
commit f497b4efd7
3 changed files with 142 additions and 61 deletions
@@ -149,9 +149,9 @@ func (m *mapSession) serveLongPoll() {
 	m.log.Trace().Caller().Msg("long poll session started")

 	// connectGen is set by [state.State.Connect] below and captured by the deferred cleanup closure.
-	// It allows [state.State.Disconnect] to reject stale calls from old sessions — if a newer session
-	// has called [state.State.Connect] (incrementing the generation), the old session's [state.State.Disconnect]
-	// sees a mismatched generation and becomes a no-op.
+	// 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
+	// the node's session count leaks and it stays online forever.
 	var connectGen uint64

 	// Clean up the session when the client disconnects
@@ -160,49 +160,55 @@ func (m *mapSession) serveLongPoll() {

 		stillConnected := m.h.mapBatcher.RemoveNode(m.node.ID, m.ch)

-		// If another session already exists for this node (reconnect
-		// happened before this cleanup ran), skip the grace period
-		// entirely — the node is not actually disconnecting.
-		if stillConnected {
+		// This session never reached [state.State.Connect]; there is no
+		// session to release.
+		if connectGen == 0 {
 			return
 		}

 		// When a node disconnects, it might rapidly reconnect (e.g. mobile clients, network weather).
 		// Instead of immediately marking the node as offline, we wait a few seconds to see if it reconnects.
-		// If it does reconnect, the existing [mapSession] will be replaced and the node remains online.
-		// If it doesn't reconnect within the timeout, we mark it as offline.
+		// If it reconnects during the wait, the new session's Connect raises the
+		// session count, so the release below keeps the node online.
 		//
 		// This avoids flapping nodes in the UI and unnecessary churn in the network.
 		// This is not my favourite solution, but it kind of works in our eventually consistent world.
-		ticker := time.NewTicker(time.Second)
-		defer ticker.Stop()
+		//
+		// When another session already replaced this one (stillConnected), skip
+		// the wait — but never the release itself. A cancelled map request whose
+		// handler ran late is exactly such a session: if it kept its session
+		// acquired on this path, the surviving session's release could never
+		// take the node offline (the relogin flake).
+		if !stillConnected {
+			// Wait up to 10 seconds for the node to reconnect.
+			// 10 seconds was arbitrary chosen as a reasonable time to reconnect.
+			ticker := time.NewTicker(time.Second)
+			defer ticker.Stop()

-		disconnected := true
-		// Wait up to 10 seconds for the node to reconnect.
-		// 10 seconds was arbitrary chosen as a reasonable time to reconnect.
-		for range 10 {
-			if m.h.mapBatcher.IsConnected(m.node.ID) {
-				disconnected = false
-				break
+			for range 10 {
+				if m.h.mapBatcher.IsConnected(m.node.ID) {
+					break
+				}
+
+				<-ticker.C
 			}
-
-			<-ticker.C
 		}

-		if disconnected {
-			// Pass the generation from our [state.State.Connect] call. If a newer session has
-			// connected since (bumping the generation), [state.State.Disconnect] will detect
-			// the mismatch and skip the state update, preventing the race where
-			// an old grace period goroutine overwrites a newer session's online status.
-			disconnectChanges, err := m.h.state.Disconnect(m.node.ID, connectGen)
-			if err != nil {
-				m.log.Error().Caller().Err(err).Msg("failed to disconnect node")
-			}
-
-			m.h.Change(disconnectChanges...)
-			m.afterServeLongPoll()
-			m.log.Info().Caller().Str(zf.Chan, fmt.Sprintf("%p", m.ch)).Msg("node has disconnected")
+		// Release this session. The node goes offline exactly when the last
+		// live session is released, so releases from replaced or stale
+		// sessions are harmless regardless of the order they run in.
+		disconnectChanges, err := m.h.state.Disconnect(m.node.ID, connectGen)
+		if err != nil {
+			m.log.Error().Caller().Err(err).Msg("failed to disconnect node")
 		}
+
+		if len(disconnectChanges) == 0 {
+			return
+		}
+
+		m.h.Change(disconnectChanges...)
+		m.afterServeLongPoll()
+		m.log.Info().Caller().Str(zf.Chan, fmt.Sprintf("%p", m.ch)).Msg("node has disconnected")
 	}()

 	// Set up the client stream
@@ -84,11 +84,10 @@ func TestConnectDisconnectOrdinaryNodeNoRuntimeRecompute(t *testing.T) {
 	})

 	t.Run("disconnect", func(t *testing.T) {
-		// Connect first so Disconnect's epoch check passes.
+		// Connect acquired a session in the connect subtest too; drain to the
+		// last release, which is the one that marks the node offline.
 		_, epoch := s.Connect(nodeID)
-
-		cs, err := s.Disconnect(nodeID, epoch)
-		require.NoError(t, err)
+		cs := drainSessions(t, s, nodeID, epoch)

 		assert.True(t, hasPeerPatch(cs),
 			"Disconnect should still emit a lightweight offline peer patch")
@@ -125,9 +124,7 @@ func TestConnectDisconnectRelayTargetTriggersRecompute(t *testing.T) {

 	t.Run("disconnect", func(t *testing.T) {
 		_, epoch := s.Connect(nodeID)
-
-		cs, err := s.Disconnect(nodeID, epoch)
-		require.NoError(t, err)
+		cs := drainSessions(t, s, nodeID, epoch)

 		assert.NotEmpty(t, runtimePeerComputationReasons(cs),
 			"relay-target node disconnect must trigger a runtime peer recompute")
@@ -158,9 +155,7 @@ func TestConnectDisconnectSubnetRouterForcesRecompute(t *testing.T) {

 	t.Run("disconnect", func(t *testing.T) {
 		_, epoch := s.Connect(nodeID)
-
-		cs, err := s.Disconnect(nodeID, epoch)
-		require.NoError(t, err)
+		cs := drainSessions(t, s, nodeID, epoch)

 		assert.True(t, forcesPeerRecompute(cs),
 			"subnet router disconnect must force a peer recompute")
@@ -207,10 +202,78 @@ func TestConnectDisconnectSubnetRouterEmitsPolicyChangeNotFull(t *testing.T) {

 	t.Run("disconnect", func(t *testing.T) {
 		_, epoch := s.Connect(nodeID)
-
-		cs, err := s.Disconnect(nodeID, epoch)
-		require.NoError(t, err)
+		cs := drainSessions(t, s, nodeID, epoch)

 		assertRecomputeNotFull(t, cs)
 	})
 }
+
+// drainSessions releases poll sessions on nodeID until the node goes
+// offline, returning the changes from the final release (the one that
+// emits the offline notifications). Fails the test if the node is
+// still online after releasing as many sessions as Connect calls could
+// plausibly have acquired.
+func drainSessions(t *testing.T, s *State, nodeID types.NodeID, epoch uint64) []change.Change {
+	t.Helper()
+
+	// Arbitrary upper bound: comfortably above the number of Connect
+	// calls any test here makes. It only guards against looping
+	// forever when the node never goes offline.
+	const maxSessions = 16
+
+	for range maxSessions {
+		cs, err := s.Disconnect(nodeID, epoch)
+		require.NoError(t, err)
+
+		if len(cs) > 0 {
+			return cs
+		}
+	}
+
+	t.Fatalf("node %d still online after releasing %d sessions", nodeID, maxSessions)
+
+	return nil
+}
+
+// TestDisconnectOutOfOrderSessionsCannotStrandNodeOnline reproduces the
+// server side of the relogin flake at the state level: a cancelled map
+// request whose handler runs late acquires a session (and the newest
+// epoch) after the real session's Connect, then releases without taking
+// the node offline because the real session is still live. The real
+// session's release — carrying the older epoch — must still take the
+// node offline. Under the old epoch-equality gate it was rejected as
+// stale and the node stayed online forever.
+func TestDisconnectOutOfOrderSessionsCannotStrandNodeOnline(t *testing.T) {
+	_, s, nodeID := persistTestSetup(t)
+	t.Cleanup(func() { _ = s.Close() })
+
+	_, liveGen := s.Connect(nodeID)
+	_, zombieGen := s.Connect(nodeID)
+	require.Greater(t, zombieGen, liveGen, "late session must hold the newer epoch")
+
+	// The zombie session dies first; another session is live, so the node
+	// must stay online and no offline changes may be emitted.
+	cs, err := s.Disconnect(nodeID, zombieGen)
+	require.NoError(t, err)
+	assert.Empty(t, cs, "release with another live session must not emit changes")
+
+	nv, ok := s.GetNodeByID(nodeID)
+	require.True(t, ok)
+
+	online, known := nv.IsOnline().GetOk()
+	require.True(t, known)
+	assert.True(t, online, "node must stay online while the real session lives")
+
+	// The real session releases last, with the older epoch. This must take
+	// the node offline.
+	cs, err = s.Disconnect(nodeID, liveGen)
+	require.NoError(t, err)
+	assert.True(t, hasPeerPatch(cs), "final release must emit the offline peer patch")
+
+	nv, ok = s.GetNodeByID(nodeID)
+	require.True(t, ok)
+
+	online, known = nv.IsOnline().GetOk()
+	require.True(t, known)
+	assert.False(t, online, "node must be offline after its last session is released")
+}
@@ -598,9 +598,9 @@ func (s *State) DeleteNode(node types.NodeView) (change.Change, error) {
 }

 // Connect marks a node connected and returns the resulting changes
-// plus a session epoch. The caller must pass the epoch back to
-// [State.Disconnect] so deferred grace-period disconnects from a
-// previous poll session are dropped (see poll.go).
+// plus a session epoch identifying this poll session. Every Connect
+// acquires one live session; the caller must release it with exactly
+// one [State.Disconnect] call once the session ends (see poll.go).
 func (s *State) Connect(id types.NodeID) ([]change.Change, uint64) {
 	prevRoutes := s.nodeStore.PrimaryRoutes()

@@ -611,6 +611,7 @@ func (s *State) Connect(id types.NodeID) ([]change.Change, uint64) {
 	node, ok := s.nodeStore.UpdateNode(id, func(n *types.Node) {
 		n.SessionEpoch++
 		epoch = n.SessionEpoch
+		n.ActiveSessions++
 		n.IsOnline = new(true)
 		n.Unhealthy = false
 	})
@@ -638,21 +639,32 @@ func (s *State) Connect(id types.NodeID) ([]change.Change, uint64) {
 	return c, epoch
 }

-// Disconnect marks the node offline. epoch must match the value
-// [State.Connect] returned for this session; otherwise the call no-ops
-// so a deferred disconnect from an older session cannot overwrite state
-// set by a newer [State.Connect]. The check and the IsOnline write share
-// an [NodeStore.UpdateNode] closure, making them atomic against
-// concurrent connects.
+// Disconnect releases one poll session previously acquired by
+// [State.Connect] and marks the node offline only when that was its
+// last live session. Sessions are counted rather than compared by
+// epoch: overlapping sessions for one node — a rapid reconnect, or a
+// cancelled map request whose handler ran late — release in any order
+// without stranding the node. An epoch-equality gate here loses when a
+// dead-on-arrival session's Connect steals the latest epoch and its
+// cleanup skips the release: the surviving session's Disconnect was
+// then rejected as stale and the node stayed online forever.
+// The count check and the IsOnline write share a
+// [NodeStore.UpdateNode] closure, making them atomic against
+// concurrent connects. epoch identifies the session for logging only.
 func (s *State) Disconnect(id types.NodeID, epoch uint64) ([]change.Change, error) {
-	var stale bool
+	var wentOffline bool

 	node, ok := s.nodeStore.UpdateNode(id, func(n *types.Node) {
-		if n.SessionEpoch != epoch {
-			stale = true
+		if n.ActiveSessions > 0 {
+			n.ActiveSessions--
+		}
+
+		if n.ActiveSessions > 0 {
 			return
 		}

+		wentOffline = true
+
 		now := time.Now()
 		n.LastSeen = &now
 		n.IsOnline = new(false)
@@ -665,11 +677,11 @@ func (s *State) Disconnect(id types.NodeID, epoch uint64) ([]change.Change, erro
 		return nil, fmt.Errorf("%w: %d", ErrNodeNotFound, id)
 	}

-	if stale {
+	if !wentOffline {
 		log.Debug().
 			Uint64("disconnect_epoch", epoch).
-			Uint64("current_epoch", node.SessionEpoch()).
-			Msg("stale disconnect rejected, newer session active")
+			Int("active_sessions", node.ActiveSessions()).
+			Msg("session released, other sessions keep node online")

 		return nil, nil
 	}