mapper: cap concurrent map sessions per node

A node chose how many map streams it held, each pinning a session.
The fifth replaces the oldest; same shape as #3468's mapper commit.

hscontrol/mapper/batcher.go

@@ -350,7 +350,10 @@ func (b *Batcher) AddNode(
 				b.totalNodes.Add(1)
 			}
 
-			existing.addConnection(newEntry)
+			existing.addConnectionWithLimit(
+				newEntry,
+				maxConcurrentMapSessionsPerNode,
+			)
 			nodeConn = existing
 
 			return existing, xsync.UpdateOp
@@ -378,6 +381,15 @@ func (b *Batcher) AddNode(
 		// cannot interleave its own lastSentPeers update.
 		nodeConn.workMu.Lock()
 
+		// Evicted by a newer stream: that stream owns the node's peer
+		// tracking and connected state, and this map may predate its own.
+		if newEntry.closed.Load() {
+			nodeConn.workMu.Unlock()
+			nlog.Debug().Caller().Msg("evicted while its initial map was in flight")
+
+			return nil
+		}
+
 		// Peers removed while the node had no stream: the initial map is a
 		// full rebuild, so follow it with the removal as a delta, see
 		// [withRemovals], and keep the removed peers tracked across the
@@ -413,6 +425,10 @@ func (b *Batcher) AddNode(
 			nodeConn.lastSentPeers.Store(peerID, struct{}{})
 		}
 
+		// Under workMu, so it lands before any newer stream's AddNode
+		// finishes and that stream can leave.
+		nodeConn.markConnected()
+
 		nodeConn.workMu.Unlock()
 
 		// Still pendingInitial, so no broadcast can race this.
@@ -431,9 +447,6 @@ func (b *Batcher) AddNode(
 		return fmt.Errorf("%w for node %d", ErrInitialMapSendTimeout, id)
 	}
 
-	// Mark the node as connected now that the initial map was sent.
-	nodeConn.markConnected()
-
 	// Node will automatically receive updates through the normal flow
 	// The initial full map already contains all current state
 

hscontrol/mapper/batcher_test.go

@@ -2048,10 +2048,17 @@ func TestBatcherMultiConnection(t *testing.T) {
 
 			thirdChannel := make(chan *tailcfg.MapResponse, 10)
 
-			err = batcher.AddNode(node1.n.ID, thirdChannel, tailcfg.CapabilityVersion(100), nil)
-			if err != nil {
-				t.Fatalf("Failed to add third connection for node1: %v", err)
-			}
+			err = batcher.Batcher.AddNode(node1.n.ID, thirdChannel, tailcfg.CapabilityVersion(100), nil)
+			require.NoError(t, err)
+
+			fourthChannel := make(chan *tailcfg.MapResponse, 10)
+			err = batcher.Batcher.AddNode(node1.n.ID, fourthChannel, tailcfg.CapabilityVersion(100), nil)
+			require.NoError(t, err)
+
+			// The fifth connection replaces the oldest while preserving the cap.
+			fifthChannel := make(chan *tailcfg.MapResponse, 10)
+			err = batcher.Batcher.AddNode(node1.n.ID, fifthChannel, tailcfg.CapabilityVersion(100), nil)
+			require.NoError(t, err)
 
 			// Yield to allow connection to be processed
 			runtime.Gosched()
@@ -2064,14 +2071,14 @@ func TestBatcherMultiConnection(t *testing.T) {
 			if info, exists := debugInfo[node1.n.ID]; exists {
 				t.Logf("Node1 debug info: %+v", info)
 
-				if info.ActiveConnections != 3 {
-					t.Errorf("Node1 should have 3 active connections, got %d", info.ActiveConnections)
+				if info.ActiveConnections != maxConcurrentMapSessionsPerNode {
+					t.Errorf("Node1 should have %d active connections, got %d", maxConcurrentMapSessionsPerNode, info.ActiveConnections)
 				} else {
-					t.Logf("SUCCESS: Node1 correctly shows 3 active connections")
+					t.Logf("SUCCESS: Node1 correctly bounds overlapping connections")
 				}
 
 				if !info.Connected {
-					t.Errorf("Node1 should show as connected with 3 active connections")
+					t.Errorf("Node1 should show as connected with overlapping connections")
 				}
 			}
 
@@ -2099,6 +2106,8 @@ func TestBatcherMultiConnection(t *testing.T) {
 			clearChannel(node1.ch)
 			clearChannel(secondChannel)
 			clearChannel(thirdChannel)
+			clearChannel(fourthChannel)
+			clearChannel(fifthChannel)
 			clearChannel(node2.ch)
 
 			// Send a change notification from node2 (so node1 should receive it on all connections)
@@ -2111,17 +2120,17 @@ func TestBatcherMultiConnection(t *testing.T) {
 				assert.Positive(c, len(node1.ch)+len(secondChannel)+len(thirdChannel), "should have received updates")
 			}, 5*time.Second, 50*time.Millisecond, "waiting for updates to propagate")
 
-			// Verify all three connections for node1 receive the update
+			// Verify the two newest connections receive the update.
 			connection1Received := false
 			connection2Received := false
 			connection3Received := false
+			connection4Received := false
+			connection5Received := false
 
 			select {
-			case mapResp := <-node1.ch:
-				connection1Received = (mapResp != nil)
-				t.Logf("Node1 connection 1 received update: %t", connection1Received)
-			case <-time.After(500 * time.Millisecond):
-				t.Errorf("Node1 connection 1 did not receive update")
+			case <-node1.ch:
+				t.Errorf("Replaced connection received an update")
+			case <-time.After(100 * time.Millisecond):
 			}
 
 			select {
@@ -2140,11 +2149,29 @@ func TestBatcherMultiConnection(t *testing.T) {
 				t.Errorf("Node1 connection 3 did not receive update")
 			}
 
-			if connection1Received && connection2Received && connection3Received {
-				t.Logf("SUCCESS: All three connections for node1 received the update")
+			select {
+			case mapResp := <-fourthChannel:
+				connection4Received = (mapResp != nil)
+				t.Logf("Node1 connection 4 received update: %t", connection4Received)
+			case <-time.After(500 * time.Millisecond):
+				t.Errorf("Node1 connection 4 did not receive update")
+			}
+
+			select {
+			case mapResp := <-fifthChannel:
+				connection5Received = (mapResp != nil)
+				t.Logf("Node1 connection 5 received update: %t", connection5Received)
+			case <-time.After(500 * time.Millisecond):
+				t.Errorf("Node1 connection 5 did not receive update")
+			}
+
+			if !connection1Received && connection2Received && connection3Received &&
+				connection4Received && connection5Received {
+				t.Logf("SUCCESS: Only the four newest connections received the update")
 			} else {
-				t.Errorf("FAILURE: Multi-connection broadcast failed - conn1: %t, conn2: %t, conn3: %t",
-					connection1Received, connection2Received, connection3Received)
+				t.Errorf("FAILURE: Bounded multi-connection broadcast failed - conn1: %t, conn2: %t, conn3: %t, conn4: %t, conn5: %t",
+					connection1Received, connection2Received, connection3Received,
+					connection4Received, connection5Received)
 			}
 
 			// Test connection removal and verify remaining connections still work.
@@ -2159,52 +2186,44 @@ func TestBatcherMultiConnection(t *testing.T) {
 			// Yield to allow removal to be processed
 			runtime.Gosched()
 
-			// Verify debug status shows 2 connections now
+			// Verify the other three admitted connections remain.
 			debugInfo2 := batcher.Debug()
 			if info, exists := debugInfo2[node1.n.ID]; exists {
-				if info.ActiveConnections != 2 {
-					t.Errorf("Node1 should have 2 active connections after removal, got %d", info.ActiveConnections)
+				if info.ActiveConnections != maxConcurrentMapSessionsPerNode-1 {
+					t.Errorf("Node1 should have %d active connections after removal, got %d",
+						maxConcurrentMapSessionsPerNode-1, info.ActiveConnections)
 				} else {
-					t.Logf("SUCCESS: Node1 correctly shows 2 active connections after removal")
+					t.Logf("SUCCESS: Node1 correctly shows %d active connections after removal",
+						maxConcurrentMapSessionsPerNode-1)
 				}
 			}
 
 			// Send another update and verify remaining connections still work
-			clearChannel(node1.ch)
 			clearChannel(thirdChannel)
 
 			testChangeSet2 := change.NodeAdded(node2.n.ID)
 
 			batcher.AddWork(testChangeSet2)
 
-			// Wait for updates to propagate to remaining channels
+			// Wait for updates to propagate to the remaining channel.
 			assert.EventuallyWithT(t, func(c *assert.CollectT) {
-				assert.Positive(c, len(node1.ch)+len(thirdChannel), "should have received updates")
+				assert.NotEmpty(c, thirdChannel, "should have received updates")
 			}, 5*time.Second, 50*time.Millisecond, "waiting for updates to propagate")
 
-			// Verify remaining connections still receive updates
+			// Verify only the remaining connection receives updates.
 			remaining1Received := false
-			remaining3Received := false
-
-			select {
-			case mapResp := <-node1.ch:
-				remaining1Received = (mapResp != nil)
-			case <-time.After(500 * time.Millisecond):
-				t.Errorf("Node1 connection 1 did not receive update after removal")
-			}
 
 			select {
 			case mapResp := <-thirdChannel:
-				remaining3Received = (mapResp != nil)
+				remaining1Received = (mapResp != nil)
 			case <-time.After(500 * time.Millisecond):
 				t.Errorf("Node1 connection 3 did not receive update after removal")
 			}
 
-			if remaining1Received && remaining3Received {
-				t.Logf("SUCCESS: Remaining connections still receive updates after removal")
+			if remaining1Received {
+				t.Logf("SUCCESS: Remaining connection still receives updates after removal")
 			} else {
-				t.Errorf("FAILURE: Remaining connections failed to receive updates - conn1: %t, conn3: %t",
-					remaining1Received, remaining3Received)
+				t.Errorf("FAILURE: Remaining connection did not receive updates")
 			}
 
 			// Drain secondChannel of any messages received before removal
@@ -2222,6 +2241,124 @@ func TestBatcherMultiConnection(t *testing.T) {
 	}
 }
 
+func TestMapSessionOverlapLimitConcurrent(t *testing.T) {
+	t.Parallel()
+
+	mc := newMultiChannelNodeConn(1, nil)
+
+	const attempts = 64
+
+	panics := runConcurrentlyWithTimeout(t, attempts, 5*time.Second, func(i int) {
+		entry := makeConnectionEntry(
+			fmt.Sprintf("connection-%d", i),
+			make(chan *tailcfg.MapResponse, 1),
+		)
+		mc.addConnectionWithLimit(entry, maxConcurrentMapSessionsPerNode)
+	})
+
+	require.Zero(t, panics)
+	assert.Equal(t, maxConcurrentMapSessionsPerNode, mc.getActiveConnectionCount())
+}
+
+// waitInitialMapSendParked waits until a goroutine is blocked in AddNode's
+// initial map send, which proves that map has been generated.
+func waitInitialMapSendParked(t *testing.T) {
+	t.Helper()
+
+	const addNode = "github.com/juanfont/headscale/hscontrol/mapper.(*Batcher).AddNode("
+
+	buf := make([]byte, 1<<22)
+
+	require.Eventually(t, func() bool {
+		stacks := string(buf[:runtime.Stack(buf, true)])
+		for g := range strings.SplitSeq(stacks, "\n\n") {
+			header, frames, _ := strings.Cut(g, "\n")
+			if strings.Contains(header, "[select") &&
+				strings.HasPrefix(frames, addNode) {
+				return true
+			}
+		}
+
+		return false
+	}, 5*time.Second, time.Millisecond)
+}
+
+// TestEvictedSessionLeavesPeerTrackingAlone evicts a stream whose initial map,
+// generated while a peer was hidden, is still in flight. The newer streams
+// list the peer, and the evicted call must not remove it from them.
+func TestEvictedSessionLeavesPeerTrackingAlone(t *testing.T) {
+	td, cleanup := setupBatcherWithTestData(
+		t,
+		NewBatcherAndMapper,
+		1,
+		2,
+		normalBufferSize,
+	)
+	defer cleanup()
+
+	b := unwrapBatcher(td.Batcher)
+	peer, target := td.Nodes[0].n.ID.NodeID(), td.Nodes[1].n.ID
+	capVer := tailcfg.CapabilityVersion(100)
+
+	peerIDs := func(resp *tailcfg.MapResponse) []tailcfg.NodeID {
+		ids := make([]tailcfg.NodeID, 0, len(resp.Peers))
+		for _, p := range resp.Peers {
+			ids = append(ids, p.ID)
+		}
+
+		return ids
+	}
+
+	_, err := td.State.SetPolicy([]byte(`{"acls":[]}`))
+	require.NoError(t, err)
+
+	// Unbuffered and unread: AddNode parks on the initial map send.
+	evictedCh := make(chan *tailcfg.MapResponse)
+	addErr := make(chan error, 1)
+
+	go func() { addErr <- b.AddNode(target, evictedCh, capVer, nil) }()
+
+	waitInitialMapSendParked(t)
+
+	_, err = td.State.SetPolicy(
+		[]byte(`{"acls":[{"action":"accept","src":["*"],"dst":["*:*"]}]}`),
+	)
+	require.NoError(t, err)
+
+	// The last of these evicts the parked stream.
+	newer := make([]chan *tailcfg.MapResponse, maxConcurrentMapSessionsPerNode)
+	for i := range newer {
+		newer[i] = make(chan *tailcfg.MapResponse, normalBufferSize)
+		require.NoError(t, b.AddNode(target, newer[i], capVer, nil))
+		require.Contains(t, peerIDs(<-newer[i]), peer, "stream %d initial map", i)
+	}
+
+	stale := <-evictedCh
+	require.NotContains(t, peerIDs(stale), peer, "the map predates the policy change")
+	require.NoError(t, <-addErr)
+
+	// Queued behind anything the evicted call queued.
+	b.AddWork(change.DERPMap())
+
+	for i, ch := range newer {
+		for done := false; !done; {
+			select {
+			case resp := <-ch:
+				require.NotContains(t, resp.PeersRemoved, peer, "stream %d", i)
+				done = resp.DERPMap != nil
+			case <-time.After(5 * time.Second):
+				t.Fatalf("stream %d never got the DERP map", i)
+			}
+		}
+	}
+
+	nc, ok := b.nodes.Load(target)
+	require.True(t, ok)
+
+	_, sent := nc.lastSentPeers.Load(peer)
+	assert.True(t, sent, "the peer stays tracked as sent")
+}
+
 // TestNodeDeletedWhileChangesPending reproduces issue #2924 where deleting a node
 // from state while there are pending changes for that node in the batcher causes
 // "node not found" errors. The race condition occurs when:

hscontrol/mapper/batcher_unit_test.go

@@ -1655,6 +1655,34 @@ func TestRemoveConnectionAtIndex_NilsTrailingSlot(t *testing.T) {
 	mc.mutex.Unlock()
 }
 
+// A node chooses how many map streams it opens and each pins a session, so
+// past the limit a new connection replaces the oldest.
+func TestAddConnectionWithLimit_StopsOldest(t *testing.T) {
+	mc := newMultiChannelNodeConn(1, nil)
+
+	var stopped []int
+
+	for i := range maxConcurrentMapSessionsPerNode + 3 {
+		entry := makeConnectionEntry(
+			fmt.Sprintf("conn-%d", i),
+			make(chan *tailcfg.MapResponse, 1),
+		)
+		entry.stop = func() { stopped = append(stopped, i) }
+		mc.addConnectionWithLimit(entry, maxConcurrentMapSessionsPerNode)
+	}
+
+	assert.Equal(t, maxConcurrentMapSessionsPerNode, mc.getActiveConnectionCount())
+	assert.Equal(t, []int{0, 1, 2}, stopped,
+		"the oldest connections should be stopped, each once")
+
+	// The newest connections are the ones kept, in arrival order.
+	mc.mutex.RLock()
+	assert.Equal(t, "conn-3", mc.connections[0].id)
+	assert.Equal(t, fmt.Sprintf("conn-%d", maxConcurrentMapSessionsPerNode+2),
+		mc.connections[len(mc.connections)-1].id)
+	mc.mutex.RUnlock()
+}
+
 // ============================================================================
 // SSHPolicy delta Tests
 // ============================================================================

hscontrol/mapper/node_conn.go

@@ -23,6 +23,8 @@ import (
 // folding to the drain.
 const minPendingFold = 64
 
+const maxConcurrentMapSessionsPerNode = 4
+
 // errNoActiveConnections is returned by [multiChannelNodeConn.send] when a node
 // has no active connections (disconnected but kept in the batcher for rapid
 // reconnection). Callers must not update peer tracking state (lastSentPeers)
@@ -225,9 +227,26 @@ func (mc *multiChannelNodeConn) removeConnectionAtIndexLocked(i int, stopConnect
 
 // addConnection adds a new connection.
 func (mc *multiChannelNodeConn) addConnection(entry *connectionEntry) {
+	mc.addConnectionWithLimit(entry, 0)
+}
+
+// addConnectionWithLimit atomically adds a connection while keeping at most
+// limit overlapping sessions. The oldest session is stopped and replaced when
+// the limit is full, allowing normal client poll churn to make progress. A zero
+// limit keeps the unbounded behavior used by low-level tests and benchmarks.
+func (mc *multiChannelNodeConn) addConnectionWithLimit(
+	entry *connectionEntry,
+	limit int,
+) {
 	mc.mutex.Lock()
 	defer mc.mutex.Unlock()
 
+	for limit > 0 && len(mc.connections) >= limit {
+		replaced := mc.removeConnectionAtIndexLocked(0, true)
+		mc.log.Debug().Str(zf.ConnID, replaced.id).
+			Msg("oldest overlapping connection replaced")
+	}
+
 	mc.connections = append(mc.connections, entry)
 	mc.log.Debug().Str(zf.ConnID, entry.id).
 		Int("total_connections", len(mc.connections)).
@@ -254,11 +273,10 @@ func (mc *multiChannelNodeConn) removeConnectionByChannel(c chan<- *tailcfg.MapR
 }
 
 // detach removes the connection for the given channel and marks the node
-// disconnected if no active connections remain.
+// disconnected if no active connections remain. A connection already gone
+// was evicted, and the newer stream that evicted it owns the node's state.
 func (mc *multiChannelNodeConn) detach(c chan<- *tailcfg.MapResponse) {
-	mc.removeConnectionByChannel(c)
-
-	if !mc.hasActiveConnections() {
+	if mc.removeConnectionByChannel(c) && !mc.hasActiveConnections() {
 		mc.markDisconnected()
 	}
 }

hscontrol/servertest/poll_race_test.go

@@ -1,8 +1,10 @@
 package servertest_test
 
 import (
+	"context"
 	"fmt"
 	"net/netip"
+	"sync"
 	"testing"
 	"time"
 
@@ -373,3 +375,85 @@ func TestPollRace(t *testing.T) {
 			"approved route should persist through reconnect")
 	})
 }
+
+// firstNetmap closes got on a stream's first netmap.
+type firstNetmap struct {
+	once *sync.Once
+	got  chan struct{}
+}
+
+func (f firstNetmap) UpdateFullNetmap(*netmap.NetworkMap) {
+	f.once.Do(func() { close(f.got) })
+}
+
+// TestMapSessionsCappedPerNode opens more map streams for one node than the
+// batcher keeps. The oldest are closed and release their sessions, the newest
+// keep streaming, and closing those takes the node offline.
+func TestMapSessionsCappedPerNode(t *testing.T) {
+	t.Parallel()
+
+	srv := servertest.NewServer(t)
+	user := srv.CreateUser(t, "streamcap-user")
+	c := servertest.NewClient(t, srv, "streamcap-node", servertest.WithUser(user))
+	c.WaitForUpdate(t, 10*time.Second)
+
+	id := findNodeID(t, srv, "streamcap-node")
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// With the client's own stream, six in all. Each is admitted before the
+	// next opens, so they age in order.
+	ended := make([]chan struct{}, 5)
+	for i := range ended {
+		ended[i] = make(chan struct{})
+		got := make(chan struct{})
+
+		go func() {
+			defer close(ended[i])
+
+			_ = c.Direct().PollNetMap(ctx, firstNetmap{once: new(sync.Once), got: got})
+		}()
+
+		select {
+		case <-got:
+		case <-ended[i]:
+		case <-time.After(10 * time.Second):
+			t.Fatalf("stream %d was never admitted", i)
+		}
+	}
+
+	require.EventuallyWithT(t, func(ct *assert.CollectT) {
+		node, ok := srv.State().GetNodeByID(id)
+		if assert.True(ct, ok) {
+			assert.Equal(ct, 4, node.ActiveSessions())
+		}
+
+		assert.Equal(ct, 4, srv.App.MapBatcher().Debug()[id].ActiveConnections)
+	}, 10*time.Second, 50*time.Millisecond, "a node keeps at most four map sessions")
+
+	select {
+	case <-ended[0]:
+	case <-time.After(10 * time.Second):
+		t.Fatal("the oldest extra stream was not closed")
+	}
+
+	for i, e := range ended[1:] {
+		select {
+		case <-e:
+			t.Fatalf("stream %d was closed, but it is among the newest four", i+1)
+		default:
+		}
+	}
+
+	cancel()
+	c.Disconnect(t)
+
+	require.EventuallyWithT(t, func(ct *assert.CollectT) {
+		node, ok := srv.State().GetNodeByID(id)
+		if assert.True(ct, ok) {
+			assert.Zero(ct, node.ActiveSessions())
+			assert.False(ct, node.IsOnline().Get())
+		}
+	}, 20*time.Second, 100*time.Millisecond, "closing every stream takes the node offline")
+}