diff --git a/hscontrol/mapper/batcher.go b/hscontrol/mapper/batcher.go index 0c19f0941..e77fc66ff 100644 --- a/hscontrol/mapper/batcher.go +++ b/hscontrol/mapper/batcher.go @@ -318,11 +318,13 @@ func (b *Batcher) AddNode( // Create new connection entry now := time.Now() newEntry := &connectionEntry{ - id: connID, - c: c, - version: version, - created: now, - stop: stop, + done: make(chan struct{}), + id: connID, + c: c, + version: version, + created: now, + stop: stop, + batcherDone: b.done, } // Block broadcast sends to this connection until its initial map // is delivered below, so a delta cannot become the stream's first @@ -632,22 +634,18 @@ func (b *Batcher) worker(workerID int) { // finish — preventing out-of-order delivery and races // on lastSentPeers (Clear+Store vs Range). if nc, exists := b.nodes.Load(w.nodeID); exists && nc != nil { - // Changes that found only connections still awaiting - // their initial map are retried next tick: the - // in-flight initial map may have been generated from - // a snapshot older than the change, so dropping them - // would lose updates. Collected and prepended as a - // group to keep their order ahead of newer pending - // changes — order matters for stateful patches like - // online/offline. + // Requeue unready or backpressured work so the next tick + // regenerates it from current state. var retry []change.Change nc.workMu.Lock() - for _, ch := range w.changes { + for i, ch := range w.changes { err := nc.change(ch) - if errors.Is(err, errNoReadyConnections) { - retry = append(retry, ch) - continue + if errors.Is(err, errNoReadyConnections) || errors.Is(err, errConnectionBackpressure) { + // Keep the unsent suffix ahead of newer pending changes. + // Older reconnect siblings can receive idempotent duplicates on retry. + retry = w.changes[i:] + break } if err != nil { @@ -664,7 +662,7 @@ func (b *Batcher) worker(workerID int) { nc.prependPending(retry...) } - // Bundle delivered; allow the next tick's bundle to queue. + // Bundle delivered or requeued; allow the next tick's bundle to queue. nc.inFlight.Store(false) } case <-b.done: diff --git a/hscontrol/mapper/batcher_bench_test.go b/hscontrol/mapper/batcher_bench_test.go index 98d9e4177..226a2d834 100644 --- a/hscontrol/mapper/batcher_bench_test.go +++ b/hscontrol/mapper/batcher_bench_test.go @@ -185,6 +185,7 @@ func benchBatcher(nodeCount, bufferSize int) (*Batcher, map[types.NodeID]chan *t mc := newMultiChannelNodeConn(id, nil) ch := make(chan *tailcfg.MapResponse, bufferSize) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("conn-%d", i), c: ch, version: tailcfg.CapabilityVersion(100), @@ -362,6 +363,7 @@ func BenchmarkMultiChannelBroadcast(b *testing.B) { for j := range 2 { ch := make(chan *tailcfg.MapResponse, b.N+1) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("extra-%d-%d", i, j), c: ch, version: tailcfg.CapabilityVersion(100), @@ -520,6 +522,7 @@ func BenchmarkConnectionChurn(b *testing.B) { // Add new connection newCh := make(chan *tailcfg.MapResponse, 10) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("churn-%d", i), c: newCh, version: tailcfg.CapabilityVersion(100), @@ -566,6 +569,7 @@ func BenchmarkConcurrentSendAndChurn(b *testing.B) { newCh := make(chan *tailcfg.MapResponse, 100) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("churn-%d", i), c: newCh, version: tailcfg.CapabilityVersion(100), diff --git a/hscontrol/mapper/batcher_concurrency_test.go b/hscontrol/mapper/batcher_concurrency_test.go index 3bdd92c43..e84ba08a8 100644 --- a/hscontrol/mapper/batcher_concurrency_test.go +++ b/hscontrol/mapper/batcher_concurrency_test.go @@ -59,6 +59,7 @@ func setupLightweightBatcher(t *testing.T, nodeCount, bufferSize int) *lightweig mc := newMultiChannelNodeConn(id, nil) // nil mapper is fine for channel tests ch := make(chan *tailcfg.MapResponse, bufferSize) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("conn-%d", i), c: ch, version: tailcfg.CapabilityVersion(100), @@ -577,11 +578,9 @@ func TestProcessBatchedChanges_ConcurrentAdd_NoDataLoss(t *testing.T) { wg.Go(func() { for i := range iterations { lb.b.addToBatch(change.Change{ - Reason: fmt.Sprintf("add-%d", i), - TargetNode: types.NodeID(1), - PeerPatches: []*tailcfg.PeerChange{ - {NodeID: tailcfg.NodeID(i + 100)}, //nolint:gosec // test - }, + Reason: fmt.Sprintf("add-%d", i), + TargetNode: types.NodeID(1), + PingRequest: &tailcfg.PingRequest{URL: fmt.Sprintf("https://ping/%d", i)}, }) addedCount.Add(1) } @@ -600,18 +599,27 @@ func TestProcessBatchedChanges_ConcurrentAdd_NoDataLoss(t *testing.T) { // One final process to flush any remaining lb.b.processBatchedChanges() - // Count the patches across all bundled work items in the channel. Each - // work item bundles a node's pending changes, coalescing peer deltas, so - // the patches, not the changes, are what must all arrive. + // Count the pings across all bundled work items in the channel: pending + // may fold or collapse state changes, but every ping must arrive. queuedChanges := 0 workItems := len(lb.b.workCh) for range workItems { w := <-lb.b.workCh - queuedChanges += countPatches(w.changes) + for _, ch := range w.changes { + if ch.PingRequest != nil { + queuedChanges++ + } + } + } + // State can coalesce to a full update, but every ping must survive. + remaining := 0 + + for _, ch := range getPendingForNode(lb.b, types.NodeID(1)) { + if ch.PingRequest != nil { + remaining++ + } } - // Also count any still-pending - remaining := countPatches(getPendingForNode(lb.b, types.NodeID(1))) total := queuedChanges + remaining added := int(addedCount.Load()) @@ -691,6 +699,7 @@ func TestWorkMu_PreventsInterTickRace(t *testing.T) { mc := newMultiChannelNodeConn(1, nil) ch := make(chan *tailcfg.MapResponse, 100) entry := &connectionEntry{ + done: make(chan struct{}), id: "test", c: ch, version: tailcfg.CapabilityVersion(100), @@ -990,6 +999,7 @@ func TestBug3_CleanupOfflineNodes_TOCTOU(t *testing.T) { newCh := make(chan *tailcfg.MapResponse, 10) entry := &connectionEntry{ + done: make(chan struct{}), id: "reconnected", c: newCh, version: tailcfg.CapabilityVersion(100), @@ -1033,6 +1043,7 @@ func TestBug3_CleanupOfflineNodes_TOCTOU(t *testing.T) { if mc, ok := lb.b.nodes.Load(targetNode); ok { reconnCh := make(chan *tailcfg.MapResponse, 10) reconnEntry := &connectionEntry{ + done: make(chan struct{}), id: "race-reconn", c: reconnCh, version: tailcfg.CapabilityVersion(100), @@ -1469,6 +1480,7 @@ func TestScale1000_MultiChannelBroadcast(t *testing.T) { ch := make(chan *tailcfg.MapResponse, bufferSize) nc.channels[j] = ch entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("conn-%d-%d", i, j), c: ch, version: tailcfg.CapabilityVersion(100), @@ -1564,6 +1576,7 @@ func TestScale1000_ConnectionChurn(t *testing.T) { // Add new connection newCh := make(chan *tailcfg.MapResponse, 20) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("churn-%d-%d", i, cycle), c: newCh, version: tailcfg.CapabilityVersion(100), @@ -1649,6 +1662,7 @@ func TestScale1000_ConcurrentAddRemove(t *testing.T) { // Add a new connection ch := make(chan *tailcfg.MapResponse, 10) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("concurrent-%d", i), c: ch, version: tailcfg.CapabilityVersion(100), @@ -1768,6 +1782,7 @@ func TestScale1000_BroadcastDuringNodeChurn(t *testing.T) { mc := newMultiChannelNodeConn(id, nil) ch := make(chan *tailcfg.MapResponse, 10) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("rechurn-%d-%d", i, cycle), c: ch, version: tailcfg.CapabilityVersion(100), @@ -1843,6 +1858,7 @@ func TestScale1000_WorkChannelSaturation(t *testing.T) { mc := newMultiChannelNodeConn(id, nil) ch := make(chan *tailcfg.MapResponse, 1) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("conn-%d", i), c: ch, version: tailcfg.CapabilityVersion(100), diff --git a/hscontrol/mapper/batcher_scale_bench_test.go b/hscontrol/mapper/batcher_scale_bench_test.go index 65e9898b3..658a8d29d 100644 --- a/hscontrol/mapper/batcher_scale_bench_test.go +++ b/hscontrol/mapper/batcher_scale_bench_test.go @@ -121,6 +121,7 @@ func BenchmarkScale_ConnectionChurn(b *testing.B) { newCh := make(chan *tailcfg.MapResponse, 10) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("sc-%d", i), c: newCh, version: tailcfg.CapabilityVersion(100), @@ -408,6 +409,7 @@ func BenchmarkScale_MultiChannelBroadcast(b *testing.B) { for j := range 2 { ch := make(chan *tailcfg.MapResponse, b.N+1) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("extra-%d-%d", i, j), c: ch, version: tailcfg.CapabilityVersion(100), @@ -513,6 +515,7 @@ func BenchmarkScale_ConcurrentSendAndChurn(b *testing.B) { newCh := make(chan *tailcfg.MapResponse, b.N+1) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("sc-churn-%d", i), c: newCh, version: tailcfg.CapabilityVersion(100), @@ -585,6 +588,7 @@ func BenchmarkScale_MixedWorkload(b *testing.B) { newCh := make(chan *tailcfg.MapResponse, 10) entry := &connectionEntry{ + done: make(chan struct{}), id: fmt.Sprintf("mix-churn-%d", i), c: newCh, version: tailcfg.CapabilityVersion(100), diff --git a/hscontrol/mapper/batcher_test.go b/hscontrol/mapper/batcher_test.go index 41b2af080..c01ef8c7e 100644 --- a/hscontrol/mapper/batcher_test.go +++ b/hscontrol/mapper/batcher_test.go @@ -1177,8 +1177,8 @@ func TestBatcherCoalescesPolicyRecomputesPerTick(t *testing.T) { // TestBatcherCoalescesPeerDeltasPerTick lands a join storm, every other node // added and coming online, in one tick. Rendered one frame per change, a -// storm outgrows the session channel within a single tick and the batcher -// drops the connection as stale; the run must reach the node as one frame. +// storm outgrows the session channel within a single tick and the rest backs +// off to later ticks; the run must reach the node as one frame. func TestBatcherCoalescesPeerDeltasPerTick(t *testing.T) { for _, bf := range allBatcherFunctions { t.Run(bf.name, func(t *testing.T) { diff --git a/hscontrol/mapper/batcher_unit_test.go b/hscontrol/mapper/batcher_unit_test.go index 991ab6780..0caa7fbdc 100644 --- a/hscontrol/mapper/batcher_unit_test.go +++ b/hscontrol/mapper/batcher_unit_test.go @@ -1,8 +1,8 @@ package mapper -// Unit tests for batcher components that do NOT require database setup. -// These tests exercise connectionEntry, multiChannelNodeConn, computePeerDiff, -// updateSentPeers, generateMapResponse branching, and handleNodeChange in isolation. +// Unit tests for batcher components: connectionEntry, multiChannelNodeConn, +// computePeerDiff, updateSentPeers, generateMapResponse branching and +// handleNodeChange. import ( "errors" @@ -11,6 +11,7 @@ import ( "sync" "sync/atomic" "testing" + "testing/synctest" "time" "github.com/juanfont/headscale/hscontrol/types" @@ -181,6 +182,7 @@ func expectNoReceive(t *testing.T, ch <-chan *tailcfg.MapResponse, timeout time. // makeConnectionEntry creates a connectionEntry with the given channel. func makeConnectionEntry(id string, ch chan<- *tailcfg.MapResponse) *connectionEntry { entry := &connectionEntry{ + done: make(chan struct{}), id: id, c: ch, version: tailcfg.CapabilityVersion(100), @@ -199,6 +201,8 @@ func TestConnectionEntry_SendSuccess(t *testing.T) { ch := make(chan *tailcfg.MapResponse, 1) entry := makeConnectionEntry("test-conn", ch) data := testMapResponse() + stalled := time.Now().Add(-connectionStallTimeout) + entry.blockedSince.Store(&stalled) beforeSend := time.Now().Unix() err := entry.send(data) @@ -206,6 +210,8 @@ func TestConnectionEntry_SendSuccess(t *testing.T) { require.NoError(t, err) assert.GreaterOrEqual(t, entry.lastUsed.Load(), beforeSend, "lastUsed should be updated after successful send") + require.ErrorIs(t, entry.send(data), errConnectionBackpressure, + "progress must reset the stall deadline before the next full queue") // Verify data was actually sent received := expectReceive(t, ch, "data should be on channel") @@ -223,18 +229,21 @@ func TestConnectionEntry_SendNilData(t *testing.T) { } func TestConnectionEntry_SendTimeout(t *testing.T) { - // Unbuffered channel with no reader = always blocks - ch := make(chan *tailcfg.MapResponse) - entry := makeConnectionEntry("test-conn", ch) - data := testMapResponse() + // Fake time keeps the bound independent of scheduler load. + synctest.Test(t, func(t *testing.T) { + // Unbuffered channel with no reader = always blocks + ch := make(chan *tailcfg.MapResponse) + entry := makeConnectionEntry("test-conn", ch) + stalled := time.Now().Add(-connectionStallTimeout) + entry.blockedSince.Store(&stalled) - start := time.Now() - err := entry.send(data) - elapsed := time.Since(start) + start := time.Now() + err := entry.send(testMapResponse()) - require.ErrorIs(t, err, ErrConnectionSendTimeout) - assert.GreaterOrEqual(t, elapsed, 40*time.Millisecond, - "should wait approximately 50ms before timeout") + require.ErrorIs(t, err, ErrConnectionSendTimeout) + assert.Less(t, time.Since(start), 100*time.Millisecond, + "stale detection must not block the worker") + }) } func TestConnectionEntry_SendClosed(t *testing.T) { @@ -266,6 +275,423 @@ func TestConnectionEntry_SendUpdatesLastUsed(t *testing.T) { "lastUsed should be updated to current time after send") } +func TestConnectionEntry_CloseInterruptsWait(t *testing.T) { + synctest.Test(t, func(t *testing.T) { + mc := newMultiChannelNodeConn(1, nil) + entry := makeConnectionEntry("blocked", make(chan *tailcfg.MapResponse)) + mc.addConnection(entry) + + result := make(chan error, 1) + go func() { result <- entry.send(testMapResponse()) }() + + synctest.Wait() + mc.close() + synctest.Wait() + + select { + case err := <-result: + require.ErrorIs(t, err, errConnectionClosed) + default: + t.Fatal("closing a connection must interrupt its send wait") + } + }) +} + +func TestConnectionEntry_WaitsForDrainingConsumer(t *testing.T) { + // Fake time: the drain always lands inside the send window. + synctest.Test(t, func(t *testing.T) { + ch := make(chan *tailcfg.MapResponse, 1) + ch <- testMapResponse() + + entry := makeConnectionEntry("draining", ch) + data := testMapResponse() + + drained := make(chan struct{}) + go func() { + defer close(drained) + + <-time.After(10 * time.Millisecond) + <-ch + }() + + err := entry.send(data) + + // Join the drainer first: fake time stops once the test body returns. + <-drained + require.NoError(t, err, "space freed within the send window must not cost another batch tick") + require.Same(t, data, expectReceive(t, ch, "send must use the freed slot")) + }) +} + +func TestBatcherRetryUsesCurrentPolicyAfterReconnect(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 2, 1) + defer cleanup() + + _, err := testData.State.SetPolicy([]byte(`{ + "acls": [{"action":"accept","src":["*"],"dst":["*:*"]}], + "ssh": [{"action":"accept","src":["autogroup:member"], + "dst":["autogroup:self"],"users":["root"]}] + }`)) + require.NoError(t, err) + + b := NewBatcher(time.Hour, 1, testData.Batcher.mapper) + + b.Start() + defer b.Close() + + id := testData.Nodes[0].n.ID + slow := make(chan *tailcfg.MapResponse, 1) + require.NoError(t, b.AddNode(id, slow, tailcfg.CurrentCapabilityVersion, nil)) + initial := expectReceive(t, slow, "permissive initial map") + require.NotEmpty(t, initial.PacketFilters["base"]) + require.NotEmpty(t, initial.SSHPolicy.Rules) + + slow <- testMapResponse() + + nc, ok := b.nodes.Load(id) + require.True(t, ok) + b.AddWork(change.PolicyOnly()) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + + _, err = testData.State.SetPolicy([]byte(`{"acls":[],"ssh":[]}`)) + require.NoError(t, err) + + fresh := make(chan *tailcfg.MapResponse, 10) + require.NoError(t, b.AddNode(id, fresh, tailcfg.CurrentCapabilityVersion, nil)) + restrictive := expectReceive(t, fresh, "restrictive initial map") + require.Empty(t, restrictive.PacketFilters["base"]) + require.Empty(t, restrictive.SSHPolicy.Rules) + require.True(t, b.RemoveNode(id, slow)) + + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + require.NotEmpty(t, fresh, "the blocked change must be retried") + + for len(fresh) > 0 { + resp := <-fresh + require.Empty(t, resp.PacketFilters["base"], "retry must never restore revoked packet access") + + if resp.SSHPolicy != nil { + require.Empty(t, resp.SSHPolicy.Rules, "retry must never restore revoked SSH access") + } + } + + require.Zero(t, b.workErrors.Load()) +} + +func TestBatcherUnreadyRetryRegeneratesRemovals(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 2, 1) + defer cleanup() + + b := NewBatcher(time.Hour, 1, testData.Batcher.mapper) + + b.Start() + defer b.Close() + + id := testData.Nodes[0].n.ID + nc := newMultiChannelNodeConn(id, b.mapper) + ch := make(chan *tailcfg.MapResponse, 10) + entry := makeConnectionEntry("initial-in-flight", ch) + entry.pendingInitial.Store(true) + nc.addConnection(entry) + b.nodes.Store(id, nc) + + initial, err := b.MapResponseFromChange(id, change.FullSelf(id)) + require.NoError(t, err) + require.Len(t, initial.Peers, 1) + peerID := initial.Peers[0].ID + + _, err = testData.State.SetPolicy([]byte(`{"acls":[]}`)) + require.NoError(t, err) + b.AddWork(change.FullUpdate()) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + require.Empty(t, ch, "broadcast must wait for initial delivery") + + // Finish the initial delivery after the worker saw empty sent-peer tracking. + nc.workMu.Lock() + ch <- initial + + nc.updateSentPeers(initial) + entry.pendingInitial.Store(false) + nc.workMu.Unlock() + require.Same(t, initial, expectReceive(t, ch, "initial delivery")) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + require.Equal(t, []tailcfg.NodeID{peerID}, expectReceive(t, ch, "regenerated removal").PeersRemoved) + require.Zero(t, b.workErrors.Load()) +} + +func TestPendingBacklogCollapsesToCurrentState(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 1, 1) + defer cleanup() + + b := NewBatcher(time.Hour, 1, testData.Batcher.mapper) + + b.Start() + defer b.Close() + + id := testData.Nodes[0].n.ID + nc := newMultiChannelNodeConn(id, b.mapper) + ch := make(chan *tailcfg.MapResponse, 10) + nc.addConnection(makeConnectionEntry("slow", ch)) + nc.lastSentPeers.Store(999, struct{}{}) + b.nodes.Store(id, nc) + + for i := range 3 { + nc.appendPending(change.PingNode(id, &tailcfg.PingRequest{URL: fmt.Sprintf("https://ping/%d", i)})) + } + + for range 2000 { + nc.appendPending(change.NodeOnline(id)) + } + + pending := nc.drainPending() + require.LessOrEqual(t, len(pending), 256, "state backlog must collapse while preserving pings") + + // A retry must also collapse against a newer full update and keep older pings first. + nc.collapsePendingToFull([]change.Change{change.PingNode(id, &tailcfg.PingRequest{URL: "https://new"})}) + nc.prependPending(pending...) + pending = nc.drainPending() + require.Len(t, pending, 5) + require.True(t, pending[0].IsFull()) + nc.prependPending(pending...) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + require.Equal(t, []tailcfg.NodeID{999}, expectReceive(t, ch, "full must retain needed removals").PeersRemoved) + require.NotNil(t, expectReceive(t, ch, "collapsed state must be a full map").Node) + + for _, url := range []string{"https://ping/0", "https://ping/1", "https://ping/2", "https://new"} { + require.Equal(t, url, expectReceive(t, ch, "pings must survive in order").PingRequest.URL) + } + + require.Empty(t, ch) + require.Zero(t, b.workErrors.Load()) +} + +func TestBatcherSlowConsumerRetriesInOrder(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 2, 1) + defer cleanup() + + b := NewBatcher(time.Hour, 1, testData.Batcher.mapper) + + b.Start() + defer b.Close() + + id := testData.Nodes[0].n.ID + nc := newMultiChannelNodeConn(id, b.mapper) + + slow := make(chan *tailcfg.MapResponse, 1) + slow <- testMapResponse() + + fast := make(chan *tailcfg.MapResponse, 100) + + nc.addConnection(makeConnectionEntry("fast", fast)) + nc.addConnection(makeConnectionEntry("slow", slow)) + nc.lastSentPeers.Store(999, struct{}{}) + b.nodes.Store(id, nc) + + otherID := testData.Nodes[1].n.ID + other := newMultiChannelNodeConn(otherID, b.mapper) + otherCh := make(chan *tailcfg.MapResponse, 1) + other.addConnection(makeConnectionEntry("other", otherCh)) + b.nodes.Store(otherID, other) + + removal := change.PeersRemoved(999) + removal.TargetNode = id + b.AddWork(removal) + + for i := range 8 { + b.AddWork(change.PingNode(id, &tailcfg.PingRequest{URL: fmt.Sprintf("https://ping/%d", i)})) + } + + b.AddWork(change.PingNode(otherID, &tailcfg.PingRequest{URL: "https://other"})) + b.processBatchedChanges() + + // Requeuing at the first blocked change costs one send window; waiting on + // each of the nine changes would take nine. + require.Eventually(t, func() bool { return !nc.inFlight.Load() && !other.inFlight.Load() }, + 250*time.Millisecond, time.Millisecond) + require.Equal(t, 2, nc.getActiveConnectionCount(), "a full queue is not a dead session") + _, tracked := nc.lastSentPeers.Load(999) + require.True(t, tracked, "partial delivery must not commit shared peer tracking") + require.Equal(t, "https://other", expectReceive(t, otherCh, "other node must progress").PingRequest.URL) + <-slow + + for i := range 9 { + // Drain slower than the old send timeout, including between retries. + <-time.After(70 * time.Millisecond) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + require.Equal(t, 2, nc.getActiveConnectionCount()) + + resp := expectReceive(t, slow, "every queued change must reach the slow session") + if i == 0 { + require.Equal(t, []tailcfg.NodeID{999}, resp.PeersRemoved) + } else { + require.NotNil(t, resp.PingRequest) + require.Equal(t, fmt.Sprintf("https://ping/%d", i-1), resp.PingRequest.URL) + } + } + + _, tracked = nc.lastSentPeers.Load(999) + require.False(t, tracked, "delivered removal must update peer tracking") + require.Empty(t, nc.drainPending()) + require.Zero(t, b.workErrors.Load()) +} + +func TestBatcherStalledConsumerEvicted(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 1, 1) + defer cleanup() + + b := NewBatcher(5*time.Millisecond, 1, testData.Batcher.mapper) + id := testData.Nodes[0].n.ID + nc := newMultiChannelNodeConn(id, b.mapper) + + ch := make(chan *tailcfg.MapResponse, 1) + ch <- testMapResponse() + + stopped := make(chan struct{}) + entry := makeConnectionEntry("stalled", ch) + entry.stop = func() { close(stopped) } + nc.addConnection(entry) + b.nodes.Store(id, nc) + + b.Start() + defer b.Close() + + b.AddWork(change.PingNode(id, &tailcfg.PingRequest{URL: "https://ping"})) + + select { + case <-stopped: + t.Fatal("transient backpressure evicted the session") + case <-time.After(100 * time.Millisecond): + } + + select { + case <-stopped: + case <-time.After(6 * time.Second): + t.Fatal("a session that never drains must be stopped") + } + + require.Eventually(t, func() bool { return nc.getActiveConnectionCount() == 0 }, time.Second, time.Millisecond) +} + +func TestBatcherCloseWithBlockedConsumers(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 1, 1) + defer cleanup() + + b := NewBatcher(time.Hour, 1, testData.Batcher.mapper) + + b.Start() + defer b.Close() + + id := testData.Nodes[0].n.ID + + var stopped atomic.Int64 + + for range 20 { + ch := make(chan *tailcfg.MapResponse, 1) + require.NoError(t, b.AddNode(id, ch, tailcfg.CurrentCapabilityVersion, func() { stopped.Add(1) })) + } + + nc, ok := b.nodes.Load(id) + require.True(t, ok) + + first := nc.connections[0] + + b.AddWork(change.PingNode(id, &tailcfg.PingRequest{URL: "https://ping"})) + b.processBatchedChanges() + require.Eventually(t, func() bool { return first.blockedSince.Load() != nil && nc.inFlight.Load() }, + time.Second, time.Millisecond) + + done := make(chan struct{}) + + go func() { + b.Close() + close(done) + }() + + select { + case <-done: + case <-time.After(200 * time.Millisecond): + t.Fatal("blocked consumers held up batcher shutdown") + } + + require.Equal(t, int64(20), stopped.Load()) +} + +func TestBatcherAlternatingDrainsRevokePolicy(t *testing.T) { + testData, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 1, 1) + defer cleanup() + + b := NewBatcher(time.Hour, 1, testData.Batcher.mapper) + + b.Start() + defer b.Close() + + id := testData.Nodes[0].n.ID + + channels := []chan *tailcfg.MapResponse{ + make(chan *tailcfg.MapResponse, 1), + make(chan *tailcfg.MapResponse, 1), + } + for _, ch := range channels { + require.NoError(t, b.AddNode(id, ch, tailcfg.CurrentCapabilityVersion, nil)) + require.NotEmpty(t, expectReceive(t, ch, "permissive initial map").PacketFilters["base"]) + + ch <- testMapResponse() + } + + nc, ok := b.nodes.Load(id) + require.True(t, ok) + + older := nc.connections[0] + _, err := testData.State.SetPolicy([]byte(`{"acls":[]}`)) + require.NoError(t, err) + b.AddWork(change.PingNode(id, &tailcfg.PingRequest{URL: "https://ping"}), change.PolicyOnly()) + + revoked := [2]bool{} + + for tick := range 8 { + side := tick % 2 + select { + case resp := <-channels[side]: + if resp.PacketFilters != nil { + require.Empty(t, resp.PacketFilters["base"], "revocation must deny access") + + revoked[side] = true + } + default: + } + + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + + if (revoked[0] || older.closed.Load()) && revoked[1] { + break + } + } + + require.True(t, revoked[0] || older.closed.Load(), "older stream must be revoked or evicted") + require.True(t, revoked[1], "alternating progress must not strand revocation behind the ping") + require.Zero(t, b.workErrors.Load()) +} + +func TestPendingFullAbsorbsRetriedPingPrefix(t *testing.T) { + nc := newMultiChannelNodeConn(1, nil) + nc.collapsePendingToFull([]change.Change{change.PingNode(1, &tailcfg.PingRequest{URL: "https://new"})}) + nc.prependPending(change.PingNode(1, &tailcfg.PingRequest{URL: "https://old"})) + nc.appendPending(change.NodeOnline(2)) + + pending := nc.drainPending() + require.Len(t, pending, 3) + require.True(t, pending[0].IsFull(), "a newer full must absorb older retried state") + require.Equal(t, "https://old", pending[1].PingRequest.URL) + require.Equal(t, "https://new", pending[2].PingRequest.URL) +} + // ============================================================================ // multiChannelNodeConn.send() Tests // ============================================================================ @@ -304,7 +730,10 @@ func TestMultiChannelSend_PartialFailure(t *testing.T) { badCh := make(chan *tailcfg.MapResponse) // unbuffered, no reader mc.addConnection(makeConnectionEntry("good-1", goodCh1)) - mc.addConnection(makeConnectionEntry("bad", badCh)) + bad := makeConnectionEntry("bad", badCh) + stalled := time.Now().Add(-connectionStallTimeout) + bad.blockedSince.Store(&stalled) + mc.addConnection(bad) mc.addConnection(makeConnectionEntry("good-2", goodCh2)) err := mc.send(testMapResponse()) @@ -324,7 +753,10 @@ func TestMultiChannelSend_AllFail(t *testing.T) { // All unbuffered channels with no readers for i := range 3 { ch := make(chan *tailcfg.MapResponse) // unbuffered - mc.addConnection(makeConnectionEntry(fmt.Sprintf("bad-%d", i), ch)) + entry := makeConnectionEntry(fmt.Sprintf("bad-%d", i), ch) + stalled := time.Now().Add(-connectionStallTimeout) + entry.blockedSince.Store(&stalled) + mc.addConnection(entry) } err := mc.send(testMapResponse()) @@ -362,7 +794,10 @@ func TestMultiChannelSend_FailedConnectionRemoved(t *testing.T) { badCh := make(chan *tailcfg.MapResponse) // unbuffered, will timeout mc.addConnection(makeConnectionEntry("good", goodCh)) - mc.addConnection(makeConnectionEntry("bad", badCh)) + bad := makeConnectionEntry("bad", badCh) + stalled := time.Now().Add(-connectionStallTimeout) + bad.blockedSince.Store(&stalled) + mc.addConnection(bad) assert.Equal(t, 2, mc.getActiveConnectionCount()) @@ -1053,7 +1488,7 @@ func TestBatcher_CloseThenStartIsNoop(t *testing.T) { b.Start() // Allow a moment for any hypothetical goroutine to appear. - time.Sleep(10 * time.Millisecond) //nolint:forbidigo // test timing + <-time.After(10 * time.Millisecond) //nolint:forbidigo // test timing goroutinesAfter := runtime.NumGoroutine() @@ -1170,7 +1605,7 @@ func TestRemoveConnectionAtIndex_NilsTrailingSlot(t *testing.T) { // Manually add three entries under the lock. entries := make([]*connectionEntry, 3) for i := range entries { - entries[i] = &connectionEntry{id: fmt.Sprintf("conn-%d", i), c: make(chan<- *tailcfg.MapResponse)} + entries[i] = &connectionEntry{done: make(chan struct{}), id: fmt.Sprintf("conn-%d", i), c: make(chan<- *tailcfg.MapResponse)} } mc.mutex.Lock() @@ -1278,6 +1713,50 @@ func TestMultiChannelSend_SSHPolicyPerConnection(t *testing.T) { assert.Same(t, data.SSHPolicy, expectReceive(t, chB, "b").SSHPolicy, "b holds an older policy") } +func TestPartialPeerAdditionThenRevocation(t *testing.T) { + td, cleanup := setupBatcherWithTestData(t, NewBatcherAndMapper, 1, 2, 1) + defer cleanup() + + _, err := td.State.SetPolicy([]byte(`{"acls":[]}`)) + require.NoError(t, err) + + b := NewBatcher(time.Hour, 1, td.Batcher.mapper) + + b.Start() + defer b.Close() + + id := td.Nodes[0].n.ID + fast := make(chan *tailcfg.MapResponse, 10) + slow := make(chan *tailcfg.MapResponse, 1) + + require.NoError(t, b.AddNode(id, fast, tailcfg.CurrentCapabilityVersion, nil)) + require.NoError(t, b.AddNode(id, slow, tailcfg.CurrentCapabilityVersion, nil)) + require.Empty(t, (<-fast).Peers) + require.Empty(t, (<-slow).Peers) + + slow <- testMapResponse() + + nc, _ := b.nodes.Load(id) + _, err = td.State.SetPolicy([]byte(`{"acls":[{"action":"accept","src":["*"],"dst":["*:*"]}]}`)) + require.NoError(t, err) + b.AddWork(change.PolicyChange()) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + accepted := expectReceive(t, fast, "partial addition") + require.Len(t, accepted.PeersChanged, 1) + peer := accepted.PeersChanged[0].ID + _, committed := nc.lastSentPeers.Load(peer) + require.False(t, committed, "backpressure must not advance confirmed peer tracking") + + _, err = td.State.SetPolicy([]byte(`{"acls":[]}`)) + require.NoError(t, err) + <-slow + b.AddWork(change.FullUpdate()) + b.processBatchedChanges() + require.Eventually(t, func() bool { return !nc.inFlight.Load() }, time.Second, time.Millisecond) + require.Equal(t, []tailcfg.NodeID{peer}, expectReceive(t, fast, "revocation must remove partially delivered peer").PeersRemoved) +} + func TestEmptyMapResponse(t *testing.T) { now := time.Now() diff --git a/hscontrol/mapper/initialmap_phantom_test.go b/hscontrol/mapper/initialmap_phantom_test.go index bd0a99108..7d0158afc 100644 --- a/hscontrol/mapper/initialmap_phantom_test.go +++ b/hscontrol/mapper/initialmap_phantom_test.go @@ -29,6 +29,7 @@ func TestUnreadyConnectionDefersBroadcastsUntilInitialMap(t *testing.T) { // in between registering the channel and delivering the initial map. nc := newMultiChannelNodeConn(targetNode.n.ID, batcher.mapper) entry := &connectionEntry{ + done: make(chan struct{}), id: "test-unready", c: targetNode.ch, version: tailcfg.CapabilityVersion(100), diff --git a/hscontrol/mapper/node_conn.go b/hscontrol/mapper/node_conn.go index fe38c3978..058d70238 100644 --- a/hscontrol/mapper/node_conn.go +++ b/hscontrol/mapper/node_conn.go @@ -39,6 +39,15 @@ var errNoActiveConnections = errors.New("no active connections") // the change, so dropping it would lose the update. var errNoReadyConnections = errors.New("no connections ready for updates") +// errConnectionBackpressure keeps unsent changes queued without declaring a busy client dead. +var errConnectionBackpressure = errors.New("connection queue full") + +// Match the initial-map timeout while staying below reconnect grace and client keepalive. +const connectionStallTimeout = 5 * time.Second + +// A full snapshot supersedes state backlog; pings remain one-shot work. +const maxPendingChanges = 256 + // connectionEntry represents a single connection to a node. type connectionEntry struct { id string // unique connection ID @@ -49,6 +58,10 @@ type connectionEntry struct { lastUsed atomic.Int64 // Unix timestamp of last successful send closed atomic.Bool // Indicates if this connection has been closed + blockedSince atomic.Pointer[time.Time] + done chan struct{} + batcherDone <-chan struct{} + // pendingInitial is set by [Batcher.AddNode] while this // connection's initial map is still in flight, and cleared once it // is delivered. Broadcast sends skip such connections so a delta @@ -142,6 +155,9 @@ type multiChannelNodeConn struct { // full peer lists (which clients interpret as "no change" when empty). // Using xsync.Map for lock-free concurrent access. lastSentPeers *xsync.Map[tailcfg.NodeID, struct{}] + + // Partial additions remain removal candidates when retries render newer state. + partialPeers *xsync.Map[tailcfg.NodeID, struct{}] } // connIDCounter is a monotonically increasing counter used to generate @@ -160,6 +176,7 @@ func newMultiChannelNodeConn(id types.NodeID, mapper *mapper) *multiChannelNodeC id: id, mapper: mapper, lastSentPeers: xsync.NewMap[tailcfg.NodeID, struct{}](), + partialPeers: xsync.NewMap[tailcfg.NodeID, struct{}](), log: log.With().Uint64(zf.NodeID, id.Uint64()).Logger(), } } @@ -179,6 +196,8 @@ func (mc *multiChannelNodeConn) close() { // at most once, even if multiple cleanup paths race to remove it. func (mc *multiChannelNodeConn) stopConnection(conn *connectionEntry) { if conn.closed.CompareAndSwap(false, true) { + close(conn.done) + if conn.stop != nil { conn.stop() } @@ -294,15 +313,22 @@ func (mc *multiChannelNodeConn) offlineDuration() time.Duration { // Thread-safe via pendingMu; does not contend with the connection mutex. func (mc *multiChannelNodeConn) appendPending(changes ...change.Change) { mc.pendingMu.Lock() + mc.pending = append(mc.pending, changes...) + fullAtHead := len(mc.pending) > 0 && mc.pending[0].IsFull() // Pending drains only once the bundle in flight is delivered, so behind // workers it grew by every change; it renders folded anyway, and - // folding twice gives what folding once does. - if len(mc.pending) >= max(mc.pendingFoldAt, minPendingFold) { + // folding twice gives what folding once does. Folding first leaves the + // cap's collapse to a full map for backlogs that do not fold. + if !fullAtHead && len(mc.pending) >= max(mc.pendingFoldAt, minPendingFold) { mc.pending = change.CoalescePeerDeltas(mc.id, change.DedupePolicyChanges(mc.pending)) mc.pendingFoldAt = 2 * len(mc.pending) } + + if fullAtHead || len(mc.pending) > maxPendingChanges { + mc.pending = change.CollapseToFull(mc.id, mc.pending) + } mc.pendingMu.Unlock() } @@ -313,7 +339,13 @@ func (mc *multiChannelNodeConn) appendPending(changes ...change.Change) { // online/offline. func (mc *multiChannelNodeConn) prependPending(changes ...change.Change) { mc.pendingMu.Lock() + full := (len(mc.pending) > 0 && mc.pending[0].IsFull()) || + (len(changes) > 0 && changes[0].IsFull()) + mc.pending = append(changes, mc.pending...) + if len(mc.pending) > maxPendingChanges || full { + mc.pending = change.CollapseToFull(mc.id, mc.pending) + } mc.pendingMu.Unlock() } @@ -357,10 +389,10 @@ func (mc *multiChannelNodeConn) takePending() []change.Change { // send broadcasts data to all active connections for the node. // -// To avoid holding the write lock during potentially slow sends (each stale -// connection can block for up to 50ms), the method snapshots connections under -// a read lock, sends without any lock held, then write-locks only to remove -// failures. New connections added between the snapshot and cleanup are safe: +// To avoid holding the write lock during bounded channel waits, the method +// snapshots connections under a read lock, sends without any lock held, then +// write-locks only to remove failures. New connections added between the +// snapshot and cleanup are safe: // they receive a full initial map via [Batcher.AddNode], so missing this update // causes no data loss. func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { @@ -406,16 +438,22 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { Msg("send: broadcasting") // Send to all connections without holding any lock. - // Stale connection timeouts (50ms each) happen here without blocking - // other goroutines that need the mutex. + // Bounded channel waits do not block other goroutines that need the mutex. var ( - lastErr error - successCount int - failed []*connectionEntry + lastErr error + successCount int + failed []*connectionEntry + blocked []*connectionEntry + newestAccepted bool ) for _, conn := range snapshot { err := conn.send(conn.withSSHPolicyDelta(conn.withSelfDelta(data))) + if errors.Is(err, errConnectionBackpressure) { + blocked = append(blocked, conn) + continue + } + if err != nil { lastErr = err @@ -426,6 +464,7 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { Msg("send: connection failed") } else { successCount++ + newestAccepted = conn == snapshot[len(snapshot)-1] if data.Node != nil { conn.lastSelf.Store(data.Node) @@ -437,6 +476,12 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { } } + // The newest ready stream supersedes blocked siblings left by reconnect overlap. + if newestAccepted { + failed = append(failed, blocked...) + blocked = nil + } + // Write-lock only to remove failed connections. if len(failed) > 0 { mc.mutex.Lock() @@ -469,6 +514,17 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { Int("failed_connections", len(failed)). Msg("send: broadcast complete") + if len(blocked) > 0 { + if successCount > 0 { + // Partial additions still need removal if a retry renders newer state without them. + for _, peer := range slices.Concat(data.Peers, data.PeersChanged) { + mc.partialPeers.Store(peer.ID, struct{}{}) + } + } + + return errConnectionBackpressure + } + // Success if at least one send succeeded if successCount > 0 { return nil @@ -477,7 +533,7 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { return fmt.Errorf("node %d: all connections failed, last error: %w", mc.id, lastErr) } -// send sends data to a single connection entry with timeout-based stale connection detection. +// send distinguishes temporary backpressure from a connection that has stopped making progress. func (entry *connectionEntry) send(data *tailcfg.MapResponse) error { if data == nil { return nil @@ -489,9 +545,7 @@ func (entry *connectionEntry) send(data *tailcfg.MapResponse) error { return fmt.Errorf("connection %s: %w", entry.id, errConnectionClosed) } - // Use a short timeout to detect stale connections where the client isn't reading the channel. - // This is critical for detecting Docker containers that are forcefully terminated - // but still have channels that appear open. + // Let a draining client free a slot before yielding this worker to other nodes. // // Use a timer rather than time.After so the timeout is explicitly released // on the fast path; both are GC-recoverable since Go 1.23. @@ -502,12 +556,28 @@ func (entry *connectionEntry) send(data *tailcfg.MapResponse) error { case entry.c <- data: // Update last used timestamp on successful send entry.lastUsed.Store(time.Now().Unix()) + entry.blockedSince.Store(nil) + return nil + case <-entry.done: + return fmt.Errorf("connection %s: %w", entry.id, errConnectionClosed) + case <-entry.batcherDone: + return ErrBatcherShuttingDown case <-timer.C: + } + + // A short stall can be a busy reader; only evict after progress stops for the stall budget. + since := entry.blockedSince.Load() + if since == nil { + now := time.Now() + entry.blockedSince.CompareAndSwap(nil, &now) + } else if time.Since(*since) >= connectionStallTimeout { // Connection is likely stale - client isn't reading from channel // This catches the case where Docker containers are killed but channels remain open return fmt.Errorf("connection %s: %w", entry.id, ErrConnectionSendTimeout) } + + return errConnectionBackpressure } // nodeID returns the node ID. @@ -538,6 +608,7 @@ func (mc *multiChannelNodeConn) updateSentPeers(resp *tailcfg.MapResponse) { // Full peer list replaces tracked state entirely if resp.Peers != nil { + mc.partialPeers.Clear() mc.lastSentPeers.Clear() for _, peer := range resp.Peers { @@ -547,11 +618,13 @@ func (mc *multiChannelNodeConn) updateSentPeers(resp *tailcfg.MapResponse) { // Incremental additions for _, peer := range resp.PeersChanged { + mc.partialPeers.Delete(peer.ID) mc.lastSentPeers.Store(peer.ID, struct{}{}) } // Incremental removals for _, id := range resp.PeersRemoved { + mc.partialPeers.Delete(id) mc.lastSentPeers.Delete(id) } } @@ -567,13 +640,16 @@ func (mc *multiChannelNodeConn) computePeerDiff(currentPeers []tailcfg.NodeID) [ var removed []tailcfg.NodeID // Find removed: in lastSentPeers but not in current - mc.lastSentPeers.Range(func(id tailcfg.NodeID, _ struct{}) bool { - if _, exists := currentSet[id]; !exists { - removed = append(removed, id) - } + for _, sent := range []*xsync.Map[tailcfg.NodeID, struct{}]{mc.lastSentPeers, mc.partialPeers} { + sent.Range(func(id tailcfg.NodeID, _ struct{}) bool { + if _, exists := currentSet[id]; !exists { + removed = append(removed, id) + currentSet[id] = struct{}{} + } - return true - }) + return true + }) + } return removed } diff --git a/hscontrol/poll_test.go b/hscontrol/poll_test.go index 4738d91c0..9dc16c46c 100644 --- a/hscontrol/poll_test.go +++ b/hscontrol/poll_test.go @@ -36,6 +36,7 @@ type delayedSuccessResponseWriter struct { mu sync.Mutex writeCount int + body []byte } // deadlineResponseWriter simulates an HTTP/2 response blocked on flow control. @@ -123,9 +124,8 @@ func (w *delayedSuccessResponseWriter) Write(data []byte) (int, error) { w.mu.Unlock() if writeCount == 1 { - // Only the first write is delayed. This simulates a transiently wedged map response: - // long enough to make the batcher time out future sends, - // but short enough that the old session can still recover if we leave it alive + // Only the first write is delayed: once it returns the session drains its + // backlog (transient) or observes the stop signal (evicted). w.firstWriteStartedOnce.Do(func() { close(w.firstWriteStarted) }) @@ -140,6 +140,10 @@ func (w *delayedSuccessResponseWriter) Write(data []byte) (int, error) { }) } + w.mu.Lock() + w.body = append(w.body, data...) + w.mu.Unlock() + return len(data), nil } @@ -160,6 +164,13 @@ func (w *delayedSuccessResponseWriter) WriteCount() int { return w.writeCount } +func (w *delayedSuccessResponseWriter) Body() string { + w.mu.Lock() + defer w.mu.Unlock() + + return string(w.body) +} + // recordingResponseWriter records the status code and whether anything was // written, so a test can tell an explicit error response apart from a handler // that returned without writing (which net/http turns into an empty 200 the @@ -467,7 +478,7 @@ func TestDeleteNodeCancelsEphemeralGC(t *testing.T) { "deleting a node must cancel its ephemeral GC timer") } -// TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession +// TestGitHubIssue3129_BlockedWriteSessionLifecycle // tests the scenario reported in // https://github.com/juanfont/headscale/issues/3129. // @@ -476,106 +487,144 @@ func TestDeleteNodeCancelsEphemeralGC(t *testing.T) { // 2. Block the first map write long enough for the session to stop draining // its buffered map-response channel. // 3. While that write is blocked, queue enough updates to fill the buffered -// channel and make the next batcher send hit the stale-send timeout. -// 4. That stale-send path removes the session from the batcher, so without an -// explicit teardown hook the old [mapSession.serveLongPoll] goroutine would stay alive -// but stop receiving future updates. -// 5. Release the blocked write and verify the batcher-side stop signal makes -// that stale session exit instead of lingering as an orphaned goroutine. -func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t *testing.T) { +// channel and make the next batcher send encounter backpressure. +// 4. A transient stall must keep the session registered and deliver the queued +// updates once the write completes. A stall past the batcher's stall budget +// must remove the connection and signal its owning session before shutdown. +// 5. After the stalled write completes, verify that [mapSession.serveLongPoll] +// exits instead of lingering alive without receiving future updates. +func TestGitHubIssue3129_BlockedWriteSessionLifecycle(t *testing.T) { t.Parallel() - app := createTestApp(t) - user := app.state.CreateUserForTest("poll-stale-session-user") - createdNode := app.state.CreateRegisteredNodeForTest(user, "poll-stale-session-node") - require.NoError(t, app.state.UpdatePolicyManagerUsersForTest()) + for _, tt := range []struct { + name string + writeDelay time.Duration + evicted bool + }{ + {name: "transient", writeDelay: 250 * time.Millisecond}, + // Well past the stall budget, so a starved runner still evicts before the write returns. + {name: "stalled", writeDelay: 10 * time.Second, evicted: true}, + } { + t.Run(tt.name, func(t *testing.T) { + t.Parallel() - app.cfg.Tuning.BatchChangeDelay = 20 * time.Millisecond - app.cfg.Tuning.NodeMapSessionBufferedChanSize = 1 + app := createTestApp(t) + user := app.state.CreateUserForTest("poll-stale-session-user") + createdNode := app.state.CreateRegisteredNodeForTest(user, "poll-stale-session-node") + require.NoError(t, app.state.UpdatePolicyManagerUsersForTest()) - app.mapBatcher.Close() - require.NoError(t, app.state.Close()) + app.cfg.Tuning.BatchChangeDelay = 20 * time.Millisecond + app.cfg.Tuning.NodeMapSessionBufferedChanSize = 1 - reloadedState, err := state.NewState(app.cfg) - require.NoError(t, err) + app.mapBatcher.Close() + require.NoError(t, app.state.Close()) - app.state = reloadedState + reloadedState, err := state.NewState(app.cfg) + require.NoError(t, err) - app.mapBatcher = mapper.NewBatcherAndMapper(app.cfg, app.state) - app.mapBatcher.Start() + app.state = reloadedState - t.Cleanup(func() { - app.mapBatcher.Close() - require.NoError(t, app.state.Close()) - }) + app.mapBatcher = mapper.NewBatcherAndMapper(app.cfg, app.state) + app.mapBatcher.Start() - nodeView, ok := app.state.GetNodeByID(createdNode.ID) - require.True(t, ok, "expected node to be present in NodeStore after reload") - require.True(t, nodeView.Valid(), "expected valid node view after reload") - node := nodeView.AsStruct() + t.Cleanup(func() { + app.mapBatcher.Close() + require.NoError(t, app.state.Close()) + }) - ctx, cancel := context.WithCancel(context.Background()) - writer := newDelayedSuccessResponseWriter(250 * time.Millisecond) - session := app.newMapSession(ctx, tailcfg.MapRequest{ - Stream: true, - Version: tailcfg.CapabilityVersion(100), - }, writer, node) + nodeView, ok := app.state.GetNodeByID(createdNode.ID) + require.True(t, ok, "expected node to be present in NodeStore after reload") + require.True(t, nodeView.Valid(), "expected valid node view after reload") + node := nodeView.AsStruct() - serveDone := make(chan struct{}) + ctx, cancel := context.WithCancel(context.Background()) + writer := newDelayedSuccessResponseWriter(tt.writeDelay) + session := app.newMapSession(ctx, tailcfg.MapRequest{ + Stream: true, + Version: tailcfg.CapabilityVersion(100), + }, writer, node) - go func() { - session.serveLongPoll() - close(serveDone) - }() + serveDone := make(chan struct{}) - t.Cleanup(func() { - cancel() - app.mapBatcher.Close() + go func() { + session.serveLongPoll() + close(serveDone) + }() - select { - case <-serveDone: - case <-time.After(2 * time.Second): - } - }) + t.Cleanup(func() { + cancel() + app.mapBatcher.Close() - select { - case <-writer.FirstWriteStarted(): - case <-time.After(2 * time.Second): - t.Fatal("expected initial map write to start") + select { + case <-serveDone: + case <-time.After(2 * time.Second): + } + }) + + select { + case <-writer.FirstWriteStarted(): + case <-time.After(2 * time.Second): + t.Fatal("expected initial map write to start") + } + + streamsClosed := make(chan struct{}) + + go func() { + app.clientStreamsOpen.Wait() + close(streamsClosed) + }() + + // Distinct pings cannot collapse into a state refresh on retry. Both must + // reach a recovering client, including the one that encountered a full channel. + const ( + firstPing = "https://example.com/queued-first" + secondPing = "https://example.com/queued-second" + ) + app.mapBatcher.AddWork( + change.PingNode(node.ID, &tailcfg.PingRequest{URL: firstPing}), + change.PingNode(node.ID, &tailcfg.PingRequest{URL: secondPing}), + ) + + if tt.evicted { + require.Eventually(t, func() bool { + return app.mapBatcher.Debug()[node.ID].ActiveConnections == 0 + }, 8*time.Second, 10*time.Millisecond, "a stalled session must be pruned from the batcher") + require.Eventually(t, session.cancelChClosed.Load, time.Second, 10*time.Millisecond, + "stale eviction must stop the session before batcher shutdown") + } + + select { + case <-writer.FirstWriteFinished(): + case <-time.After(tt.writeDelay + time.Second): + t.Fatal("expected the blocked write to eventually complete") + } + + if !tt.evicted { + require.EventuallyWithT(t, func(c *assert.CollectT) { + body := writer.Body() + assert.Contains(c, body, firstPing) + assert.Contains(c, body, secondPing) + }, 2*time.Second, 10*time.Millisecond, "a transiently blocked session must receive both queued pings") + assert.Equal(t, 1, app.mapBatcher.Debug()[node.ID].ActiveConnections, + "transient backpressure must not evict the session") + assert.False(t, session.cancelChClosed.Load(), "transient backpressure must not stop the session") + + return + } + + // Shutdown skips reconnect grace; the stream wait includes disconnect cleanup. + app.mapBatcher.Close() + + assert.Eventually(t, func() bool { + select { + case <-streamsClosed: + return true + default: + return false + } + }, time.Second, 20*time.Millisecond, "after stale-send cleanup, the stale session should exit") + }) } - - streamsClosed := make(chan struct{}) - - go func() { - app.clientStreamsOpen.Wait() - close(streamsClosed) - }() - - // One update fills the buffered session channel while the first write is blocked. - // The second update then hits the 50ms stale-send timeout, so the batcher prunes - // the stale connection and triggers its stop hook. - app.mapBatcher.AddWork(change.SelfUpdate(node.ID), change.SelfUpdate(node.ID)) - - select { - case <-writer.FirstWriteFinished(): - case <-time.After(2 * time.Second): - 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: - return true - default: - return false - } - }, time.Second, 20*time.Millisecond, "after stale-send cleanup, the stale session should exit") } // TestDeletedNodeEndsLongPoll proves that deleting a node ends its long-poll