diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go index e77fc66ff..526bc2e03 100644 --- a/hscontrol/mapper/batcher.go +++ b/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 diff --git a/hscontrol/mapper/batcher_test.go b/hscontrol/mapper/batcher_test.go index fa59e1a2e..5923cb5fe 100644 --- a/hscontrol/mapper/batcher_test.go +++ b/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: diff --git a/hscontrol/mapper/batcher_unit_test.go b/hscontrol/mapper/batcher_unit_test.go index f50b08574..85f52ef43 100644 --- a/hscontrol/mapper/batcher_unit_test.go +++ b/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 // ============================================================================ diff --git a/hscontrol/mapper/node_conn.go b/hscontrol/mapper/node_conn.go index 8fb15f716..5af1399c0 100644 --- a/hscontrol/mapper/node_conn.go +++ b/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() } } diff --git a/hscontrol/servertest/poll_race_test.go b/hscontrol/servertest/poll_race_test.go index aea59e5ce..beaf6fc42 100644 --- a/hscontrol/servertest/poll_race_test.go +++ b/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") +}