mapper: retry slow sessions from current state

A full session channel after the send window evicted live but busy clients.
Requeue and regenerate; evict after no progress for connectionStallTimeout.
Retire blocked reconnect siblings once the newest stream accepts the frame.

Updates #3129.
This commit is contained in:
Kristoffer Dalby committed 2026-10-23 18:23:53 +02:00
1 parent f46ea3d1df
commit daad12e964
9 files changed
+697 -70

No files matched your search

+11 -13
View File
@@ -318,11 +318,13 @@ func (b *Batcher) AddNode(
// Create new connection entry // Create new connection entry
now := time.Now() now := time.Now()
newEntry := &connectionEntry{ newEntry := &connectionEntry{
done: make(chan struct{}),
id: connID, id: connID,
c: c, c: c,
version: version, version: version,
created: now, created: now,
stop: stop, stop: stop,
batcherDone: b.done,
} }
// Block broadcast sends to this connection until its initial map // Block broadcast sends to this connection until its initial map
// is delivered below, so a delta cannot become the stream's first // 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 // finish — preventing out-of-order delivery and races
// on lastSentPeers (Clear+Store vs Range). // on lastSentPeers (Clear+Store vs Range).
if nc, exists := b.nodes.Load(w.nodeID); exists && nc != nil { if nc, exists := b.nodes.Load(w.nodeID); exists && nc != nil {
// Changes that found only connections still awaiting // Requeue unready or backpressured work so the next tick
// their initial map are retried next tick: the // regenerates it from current state.
// 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.
var retry []change.Change var retry []change.Change
nc.workMu.Lock() nc.workMu.Lock()
for _, ch := range w.changes { for i, ch := range w.changes {
err := nc.change(ch) err := nc.change(ch)
if errors.Is(err, errNoReadyConnections) { if errors.Is(err, errNoReadyConnections) || errors.Is(err, errConnectionBackpressure) {
retry = append(retry, ch) // Keep the unsent suffix ahead of newer pending changes.
continue // Older reconnect siblings can receive idempotent duplicates on retry.
retry = w.changes[i:]
break
} }
if err != nil { if err != nil {
@@ -664,7 +662,7 @@ func (b *Batcher) worker(workerID int) {
nc.prependPending(retry...) 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) nc.inFlight.Store(false)
} }
case <-b.done: case <-b.done:
+4
View File
@@ -185,6 +185,7 @@ func benchBatcher(nodeCount, bufferSize int) (*Batcher, map[types.NodeID]chan *t
mc := newMultiChannelNodeConn(id, nil) mc := newMultiChannelNodeConn(id, nil)
ch := make(chan *tailcfg.MapResponse, bufferSize) ch := make(chan *tailcfg.MapResponse, bufferSize)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("conn-%d", i), id: fmt.Sprintf("conn-%d", i),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -362,6 +363,7 @@ func BenchmarkMultiChannelBroadcast(b *testing.B) {
for j := range 2 { for j := range 2 {
ch := make(chan *tailcfg.MapResponse, b.N+1) ch := make(chan *tailcfg.MapResponse, b.N+1)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("extra-%d-%d", i, j), id: fmt.Sprintf("extra-%d-%d", i, j),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -520,6 +522,7 @@ func BenchmarkConnectionChurn(b *testing.B) {
// Add new connection // Add new connection
newCh := make(chan *tailcfg.MapResponse, 10) newCh := make(chan *tailcfg.MapResponse, 10)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("churn-%d", i), id: fmt.Sprintf("churn-%d", i),
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -566,6 +569,7 @@ func BenchmarkConcurrentSendAndChurn(b *testing.B) {
newCh := make(chan *tailcfg.MapResponse, 100) newCh := make(chan *tailcfg.MapResponse, 100)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("churn-%d", i), id: fmt.Sprintf("churn-%d", i),
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
+25 -9
View File
@@ -59,6 +59,7 @@ func setupLightweightBatcher(t *testing.T, nodeCount, bufferSize int) *lightweig
mc := newMultiChannelNodeConn(id, nil) // nil mapper is fine for channel tests mc := newMultiChannelNodeConn(id, nil) // nil mapper is fine for channel tests
ch := make(chan *tailcfg.MapResponse, bufferSize) ch := make(chan *tailcfg.MapResponse, bufferSize)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("conn-%d", i), id: fmt.Sprintf("conn-%d", i),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -579,9 +580,7 @@ func TestProcessBatchedChanges_ConcurrentAdd_NoDataLoss(t *testing.T) {
lb.b.addToBatch(change.Change{ lb.b.addToBatch(change.Change{
Reason: fmt.Sprintf("add-%d", i), Reason: fmt.Sprintf("add-%d", i),
TargetNode: types.NodeID(1), TargetNode: types.NodeID(1),
PeerPatches: []*tailcfg.PeerChange{ PingRequest: &tailcfg.PingRequest{URL: fmt.Sprintf("https://ping/%d", i)},
{NodeID: tailcfg.NodeID(i + 100)}, //nolint:gosec // test
},
}) })
addedCount.Add(1) addedCount.Add(1)
} }
@@ -600,18 +599,27 @@ func TestProcessBatchedChanges_ConcurrentAdd_NoDataLoss(t *testing.T) {
// One final process to flush any remaining // One final process to flush any remaining
lb.b.processBatchedChanges() lb.b.processBatchedChanges()
// Count the patches across all bundled work items in the channel. Each // Count the pings across all bundled work items in the channel: pending
// work item bundles a node's pending changes, coalescing peer deltas, so // may fold or collapse state changes, but every ping must arrive.
// the patches, not the changes, are what must all arrive.
queuedChanges := 0 queuedChanges := 0
workItems := len(lb.b.workCh) workItems := len(lb.b.workCh)
for range workItems { for range workItems {
w := <-lb.b.workCh 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 total := queuedChanges + remaining
added := int(addedCount.Load()) added := int(addedCount.Load())
@@ -691,6 +699,7 @@ func TestWorkMu_PreventsInterTickRace(t *testing.T) {
mc := newMultiChannelNodeConn(1, nil) mc := newMultiChannelNodeConn(1, nil)
ch := make(chan *tailcfg.MapResponse, 100) ch := make(chan *tailcfg.MapResponse, 100)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: "test", id: "test",
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -990,6 +999,7 @@ func TestBug3_CleanupOfflineNodes_TOCTOU(t *testing.T) {
newCh := make(chan *tailcfg.MapResponse, 10) newCh := make(chan *tailcfg.MapResponse, 10)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: "reconnected", id: "reconnected",
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -1033,6 +1043,7 @@ func TestBug3_CleanupOfflineNodes_TOCTOU(t *testing.T) {
if mc, ok := lb.b.nodes.Load(targetNode); ok { if mc, ok := lb.b.nodes.Load(targetNode); ok {
reconnCh := make(chan *tailcfg.MapResponse, 10) reconnCh := make(chan *tailcfg.MapResponse, 10)
reconnEntry := &connectionEntry{ reconnEntry := &connectionEntry{
done: make(chan struct{}),
id: "race-reconn", id: "race-reconn",
c: reconnCh, c: reconnCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -1469,6 +1480,7 @@ func TestScale1000_MultiChannelBroadcast(t *testing.T) {
ch := make(chan *tailcfg.MapResponse, bufferSize) ch := make(chan *tailcfg.MapResponse, bufferSize)
nc.channels[j] = ch nc.channels[j] = ch
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("conn-%d-%d", i, j), id: fmt.Sprintf("conn-%d-%d", i, j),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -1564,6 +1576,7 @@ func TestScale1000_ConnectionChurn(t *testing.T) {
// Add new connection // Add new connection
newCh := make(chan *tailcfg.MapResponse, 20) newCh := make(chan *tailcfg.MapResponse, 20)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("churn-%d-%d", i, cycle), id: fmt.Sprintf("churn-%d-%d", i, cycle),
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -1649,6 +1662,7 @@ func TestScale1000_ConcurrentAddRemove(t *testing.T) {
// Add a new connection // Add a new connection
ch := make(chan *tailcfg.MapResponse, 10) ch := make(chan *tailcfg.MapResponse, 10)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("concurrent-%d", i), id: fmt.Sprintf("concurrent-%d", i),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -1768,6 +1782,7 @@ func TestScale1000_BroadcastDuringNodeChurn(t *testing.T) {
mc := newMultiChannelNodeConn(id, nil) mc := newMultiChannelNodeConn(id, nil)
ch := make(chan *tailcfg.MapResponse, 10) ch := make(chan *tailcfg.MapResponse, 10)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("rechurn-%d-%d", i, cycle), id: fmt.Sprintf("rechurn-%d-%d", i, cycle),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -1843,6 +1858,7 @@ func TestScale1000_WorkChannelSaturation(t *testing.T) {
mc := newMultiChannelNodeConn(id, nil) mc := newMultiChannelNodeConn(id, nil)
ch := make(chan *tailcfg.MapResponse, 1) ch := make(chan *tailcfg.MapResponse, 1)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("conn-%d", i), id: fmt.Sprintf("conn-%d", i),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -121,6 +121,7 @@ func BenchmarkScale_ConnectionChurn(b *testing.B) {
newCh := make(chan *tailcfg.MapResponse, 10) newCh := make(chan *tailcfg.MapResponse, 10)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("sc-%d", i), id: fmt.Sprintf("sc-%d", i),
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -408,6 +409,7 @@ func BenchmarkScale_MultiChannelBroadcast(b *testing.B) {
for j := range 2 { for j := range 2 {
ch := make(chan *tailcfg.MapResponse, b.N+1) ch := make(chan *tailcfg.MapResponse, b.N+1)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("extra-%d-%d", i, j), id: fmt.Sprintf("extra-%d-%d", i, j),
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -513,6 +515,7 @@ func BenchmarkScale_ConcurrentSendAndChurn(b *testing.B) {
newCh := make(chan *tailcfg.MapResponse, b.N+1) newCh := make(chan *tailcfg.MapResponse, b.N+1)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("sc-churn-%d", i), id: fmt.Sprintf("sc-churn-%d", i),
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -585,6 +588,7 @@ func BenchmarkScale_MixedWorkload(b *testing.B) {
newCh := make(chan *tailcfg.MapResponse, 10) newCh := make(chan *tailcfg.MapResponse, 10)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: fmt.Sprintf("mix-churn-%d", i), id: fmt.Sprintf("mix-churn-%d", i),
c: newCh, c: newCh,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
+2 -2
View File
@@ -1177,8 +1177,8 @@ func TestBatcherCoalescesPolicyRecomputesPerTick(t *testing.T) {
// TestBatcherCoalescesPeerDeltasPerTick lands a join storm, every other node // TestBatcherCoalescesPeerDeltasPerTick lands a join storm, every other node
// added and coming online, in one tick. Rendered one frame per change, a // 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 // storm outgrows the session channel within a single tick and the rest backs
// drops the connection as stale; the run must reach the node as one frame. // off to later ticks; the run must reach the node as one frame.
func TestBatcherCoalescesPeerDeltasPerTick(t *testing.T) { func TestBatcherCoalescesPeerDeltasPerTick(t *testing.T) {
for _, bf := range allBatcherFunctions { for _, bf := range allBatcherFunctions {
t.Run(bf.name, func(t *testing.T) { t.Run(bf.name, func(t *testing.T) {
+492 -13
View File
@@ -1,8 +1,8 @@
package mapper package mapper
// Unit tests for batcher components that do NOT require database setup. // Unit tests for batcher components: connectionEntry, multiChannelNodeConn,
// These tests exercise connectionEntry, multiChannelNodeConn, computePeerDiff, // computePeerDiff, updateSentPeers, generateMapResponse branching and
// updateSentPeers, generateMapResponse branching, and handleNodeChange in isolation. // handleNodeChange.
import ( import (
"errors" "errors"
@@ -11,6 +11,7 @@ import (
"sync" "sync"
"sync/atomic" "sync/atomic"
"testing" "testing"
"testing/synctest"
"time" "time"
"github.com/juanfont/headscale/hscontrol/types" "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. // makeConnectionEntry creates a connectionEntry with the given channel.
func makeConnectionEntry(id string, ch chan<- *tailcfg.MapResponse) *connectionEntry { func makeConnectionEntry(id string, ch chan<- *tailcfg.MapResponse) *connectionEntry {
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: id, id: id,
c: ch, c: ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
@@ -199,6 +201,8 @@ func TestConnectionEntry_SendSuccess(t *testing.T) {
ch := make(chan *tailcfg.MapResponse, 1) ch := make(chan *tailcfg.MapResponse, 1)
entry := makeConnectionEntry("test-conn", ch) entry := makeConnectionEntry("test-conn", ch)
data := testMapResponse() data := testMapResponse()
stalled := time.Now().Add(-connectionStallTimeout)
entry.blockedSince.Store(&stalled)
beforeSend := time.Now().Unix() beforeSend := time.Now().Unix()
err := entry.send(data) err := entry.send(data)
@@ -206,6 +210,8 @@ func TestConnectionEntry_SendSuccess(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
assert.GreaterOrEqual(t, entry.lastUsed.Load(), beforeSend, assert.GreaterOrEqual(t, entry.lastUsed.Load(), beforeSend,
"lastUsed should be updated after successful send") "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 // Verify data was actually sent
received := expectReceive(t, ch, "data should be on channel") 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) { func TestConnectionEntry_SendTimeout(t *testing.T) {
// Fake time keeps the bound independent of scheduler load.
synctest.Test(t, func(t *testing.T) {
// Unbuffered channel with no reader = always blocks // Unbuffered channel with no reader = always blocks
ch := make(chan *tailcfg.MapResponse) ch := make(chan *tailcfg.MapResponse)
entry := makeConnectionEntry("test-conn", ch) entry := makeConnectionEntry("test-conn", ch)
data := testMapResponse() stalled := time.Now().Add(-connectionStallTimeout)
entry.blockedSince.Store(&stalled)
start := time.Now() start := time.Now()
err := entry.send(data) err := entry.send(testMapResponse())
elapsed := time.Since(start)
require.ErrorIs(t, err, ErrConnectionSendTimeout) require.ErrorIs(t, err, ErrConnectionSendTimeout)
assert.GreaterOrEqual(t, elapsed, 40*time.Millisecond, assert.Less(t, time.Since(start), 100*time.Millisecond,
"should wait approximately 50ms before timeout") "stale detection must not block the worker")
})
} }
func TestConnectionEntry_SendClosed(t *testing.T) { 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") "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 // multiChannelNodeConn.send() Tests
// ============================================================================ // ============================================================================
@@ -304,7 +730,10 @@ func TestMultiChannelSend_PartialFailure(t *testing.T) {
badCh := make(chan *tailcfg.MapResponse) // unbuffered, no reader badCh := make(chan *tailcfg.MapResponse) // unbuffered, no reader
mc.addConnection(makeConnectionEntry("good-1", goodCh1)) 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)) mc.addConnection(makeConnectionEntry("good-2", goodCh2))
err := mc.send(testMapResponse()) err := mc.send(testMapResponse())
@@ -324,7 +753,10 @@ func TestMultiChannelSend_AllFail(t *testing.T) {
// All unbuffered channels with no readers // All unbuffered channels with no readers
for i := range 3 { for i := range 3 {
ch := make(chan *tailcfg.MapResponse) // unbuffered 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()) err := mc.send(testMapResponse())
@@ -362,7 +794,10 @@ func TestMultiChannelSend_FailedConnectionRemoved(t *testing.T) {
badCh := make(chan *tailcfg.MapResponse) // unbuffered, will timeout badCh := make(chan *tailcfg.MapResponse) // unbuffered, will timeout
mc.addConnection(makeConnectionEntry("good", goodCh)) 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()) assert.Equal(t, 2, mc.getActiveConnectionCount())
@@ -1053,7 +1488,7 @@ func TestBatcher_CloseThenStartIsNoop(t *testing.T) {
b.Start() b.Start()
// Allow a moment for any hypothetical goroutine to appear. // 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() goroutinesAfter := runtime.NumGoroutine()
@@ -1170,7 +1605,7 @@ func TestRemoveConnectionAtIndex_NilsTrailingSlot(t *testing.T) {
// Manually add three entries under the lock. // Manually add three entries under the lock.
entries := make([]*connectionEntry, 3) entries := make([]*connectionEntry, 3)
for i := range entries { 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() 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") 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) { func TestEmptyMapResponse(t *testing.T) {
now := time.Now() now := time.Now()
@@ -29,6 +29,7 @@ func TestUnreadyConnectionDefersBroadcastsUntilInitialMap(t *testing.T) {
// in between registering the channel and delivering the initial map. // in between registering the channel and delivering the initial map.
nc := newMultiChannelNodeConn(targetNode.n.ID, batcher.mapper) nc := newMultiChannelNodeConn(targetNode.n.ID, batcher.mapper)
entry := &connectionEntry{ entry := &connectionEntry{
done: make(chan struct{}),
id: "test-unready", id: "test-unready",
c: targetNode.ch, c: targetNode.ch,
version: tailcfg.CapabilityVersion(100), version: tailcfg.CapabilityVersion(100),
+89 -13
View File
@@ -39,6 +39,15 @@ var errNoActiveConnections = errors.New("no active connections")
// the change, so dropping it would lose the update. // the change, so dropping it would lose the update.
var errNoReadyConnections = errors.New("no connections ready for updates") 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. // connectionEntry represents a single connection to a node.
type connectionEntry struct { type connectionEntry struct {
id string // unique connection ID id string // unique connection ID
@@ -49,6 +58,10 @@ type connectionEntry struct {
lastUsed atomic.Int64 // Unix timestamp of last successful send lastUsed atomic.Int64 // Unix timestamp of last successful send
closed atomic.Bool // Indicates if this connection has been closed 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 // pendingInitial is set by [Batcher.AddNode] while this
// connection's initial map is still in flight, and cleared once it // connection's initial map is still in flight, and cleared once it
// is delivered. Broadcast sends skip such connections so a delta // 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). // full peer lists (which clients interpret as "no change" when empty).
// Using xsync.Map for lock-free concurrent access. // Using xsync.Map for lock-free concurrent access.
lastSentPeers *xsync.Map[tailcfg.NodeID, struct{}] 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 // connIDCounter is a monotonically increasing counter used to generate
@@ -160,6 +176,7 @@ func newMultiChannelNodeConn(id types.NodeID, mapper *mapper) *multiChannelNodeC
id: id, id: id,
mapper: mapper, mapper: mapper,
lastSentPeers: xsync.NewMap[tailcfg.NodeID, struct{}](), lastSentPeers: xsync.NewMap[tailcfg.NodeID, struct{}](),
partialPeers: xsync.NewMap[tailcfg.NodeID, struct{}](),
log: log.With().Uint64(zf.NodeID, id.Uint64()).Logger(), 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. // at most once, even if multiple cleanup paths race to remove it.
func (mc *multiChannelNodeConn) stopConnection(conn *connectionEntry) { func (mc *multiChannelNodeConn) stopConnection(conn *connectionEntry) {
if conn.closed.CompareAndSwap(false, true) { if conn.closed.CompareAndSwap(false, true) {
close(conn.done)
if conn.stop != nil { if conn.stop != nil {
conn.stop() conn.stop()
} }
@@ -294,15 +313,22 @@ func (mc *multiChannelNodeConn) offlineDuration() time.Duration {
// Thread-safe via pendingMu; does not contend with the connection mutex. // Thread-safe via pendingMu; does not contend with the connection mutex.
func (mc *multiChannelNodeConn) appendPending(changes ...change.Change) { func (mc *multiChannelNodeConn) appendPending(changes ...change.Change) {
mc.pendingMu.Lock() mc.pendingMu.Lock()
mc.pending = append(mc.pending, changes...) 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 // Pending drains only once the bundle in flight is delivered, so behind
// workers it grew by every change; it renders folded anyway, and // workers it grew by every change; it renders folded anyway, and
// folding twice gives what folding once does. // folding twice gives what folding once does. Folding first leaves the
if len(mc.pending) >= max(mc.pendingFoldAt, minPendingFold) { // 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.pending = change.CoalescePeerDeltas(mc.id, change.DedupePolicyChanges(mc.pending))
mc.pendingFoldAt = 2 * len(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() mc.pendingMu.Unlock()
} }
@@ -313,7 +339,13 @@ func (mc *multiChannelNodeConn) appendPending(changes ...change.Change) {
// online/offline. // online/offline.
func (mc *multiChannelNodeConn) prependPending(changes ...change.Change) { func (mc *multiChannelNodeConn) prependPending(changes ...change.Change) {
mc.pendingMu.Lock() mc.pendingMu.Lock()
full := (len(mc.pending) > 0 && mc.pending[0].IsFull()) ||
(len(changes) > 0 && changes[0].IsFull())
mc.pending = append(changes, mc.pending...) mc.pending = append(changes, mc.pending...)
if len(mc.pending) > maxPendingChanges || full {
mc.pending = change.CollapseToFull(mc.id, mc.pending)
}
mc.pendingMu.Unlock() mc.pendingMu.Unlock()
} }
@@ -357,10 +389,10 @@ func (mc *multiChannelNodeConn) takePending() []change.Change {
// send broadcasts data to all active connections for the node. // send broadcasts data to all active connections for the node.
// //
// To avoid holding the write lock during potentially slow sends (each stale // To avoid holding the write lock during bounded channel waits, the method
// connection can block for up to 50ms), the method snapshots connections under // snapshots connections under a read lock, sends without any lock held, then
// a read lock, sends without any lock held, then write-locks only to remove // write-locks only to remove failures. New connections added between the
// failures. New connections added between the snapshot and cleanup are safe: // snapshot and cleanup are safe:
// they receive a full initial map via [Batcher.AddNode], so missing this update // they receive a full initial map via [Batcher.AddNode], so missing this update
// causes no data loss. // causes no data loss.
func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error { func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error {
@@ -406,16 +438,22 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error {
Msg("send: broadcasting") Msg("send: broadcasting")
// Send to all connections without holding any lock. // Send to all connections without holding any lock.
// Stale connection timeouts (50ms each) happen here without blocking // Bounded channel waits do not block other goroutines that need the mutex.
// other goroutines that need the mutex.
var ( var (
lastErr error lastErr error
successCount int successCount int
failed []*connectionEntry failed []*connectionEntry
blocked []*connectionEntry
newestAccepted bool
) )
for _, conn := range snapshot { for _, conn := range snapshot {
err := conn.send(conn.withSSHPolicyDelta(conn.withSelfDelta(data))) err := conn.send(conn.withSSHPolicyDelta(conn.withSelfDelta(data)))
if errors.Is(err, errConnectionBackpressure) {
blocked = append(blocked, conn)
continue
}
if err != nil { if err != nil {
lastErr = err lastErr = err
@@ -426,6 +464,7 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error {
Msg("send: connection failed") Msg("send: connection failed")
} else { } else {
successCount++ successCount++
newestAccepted = conn == snapshot[len(snapshot)-1]
if data.Node != nil { if data.Node != nil {
conn.lastSelf.Store(data.Node) 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. // Write-lock only to remove failed connections.
if len(failed) > 0 { if len(failed) > 0 {
mc.mutex.Lock() mc.mutex.Lock()
@@ -469,6 +514,17 @@ func (mc *multiChannelNodeConn) send(data *tailcfg.MapResponse) error {
Int("failed_connections", len(failed)). Int("failed_connections", len(failed)).
Msg("send: broadcast complete") 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 // Success if at least one send succeeded
if successCount > 0 { if successCount > 0 {
return nil 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) 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 { func (entry *connectionEntry) send(data *tailcfg.MapResponse) error {
if data == nil { if data == nil {
return nil return nil
@@ -489,9 +545,7 @@ func (entry *connectionEntry) send(data *tailcfg.MapResponse) error {
return fmt.Errorf("connection %s: %w", entry.id, errConnectionClosed) 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. // Let a draining client free a slot before yielding this worker to other nodes.
// This is critical for detecting Docker containers that are forcefully terminated
// but still have channels that appear open.
// //
// Use a timer rather than time.After so the timeout is explicitly released // 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. // 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: case entry.c <- data:
// Update last used timestamp on successful send // Update last used timestamp on successful send
entry.lastUsed.Store(time.Now().Unix()) entry.lastUsed.Store(time.Now().Unix())
entry.blockedSince.Store(nil)
return nil return nil
case <-entry.done:
return fmt.Errorf("connection %s: %w", entry.id, errConnectionClosed)
case <-entry.batcherDone:
return ErrBatcherShuttingDown
case <-timer.C: 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 // Connection is likely stale - client isn't reading from channel
// This catches the case where Docker containers are killed but channels remain open // This catches the case where Docker containers are killed but channels remain open
return fmt.Errorf("connection %s: %w", entry.id, ErrConnectionSendTimeout) return fmt.Errorf("connection %s: %w", entry.id, ErrConnectionSendTimeout)
} }
return errConnectionBackpressure
} }
// nodeID returns the node ID. // nodeID returns the node ID.
@@ -538,6 +608,7 @@ func (mc *multiChannelNodeConn) updateSentPeers(resp *tailcfg.MapResponse) {
// Full peer list replaces tracked state entirely // Full peer list replaces tracked state entirely
if resp.Peers != nil { if resp.Peers != nil {
mc.partialPeers.Clear()
mc.lastSentPeers.Clear() mc.lastSentPeers.Clear()
for _, peer := range resp.Peers { for _, peer := range resp.Peers {
@@ -547,11 +618,13 @@ func (mc *multiChannelNodeConn) updateSentPeers(resp *tailcfg.MapResponse) {
// Incremental additions // Incremental additions
for _, peer := range resp.PeersChanged { for _, peer := range resp.PeersChanged {
mc.partialPeers.Delete(peer.ID)
mc.lastSentPeers.Store(peer.ID, struct{}{}) mc.lastSentPeers.Store(peer.ID, struct{}{})
} }
// Incremental removals // Incremental removals
for _, id := range resp.PeersRemoved { for _, id := range resp.PeersRemoved {
mc.partialPeers.Delete(id)
mc.lastSentPeers.Delete(id) mc.lastSentPeers.Delete(id)
} }
} }
@@ -567,13 +640,16 @@ func (mc *multiChannelNodeConn) computePeerDiff(currentPeers []tailcfg.NodeID) [
var removed []tailcfg.NodeID var removed []tailcfg.NodeID
// Find removed: in lastSentPeers but not in current // Find removed: in lastSentPeers but not in current
mc.lastSentPeers.Range(func(id tailcfg.NodeID, _ struct{}) bool { 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 { if _, exists := currentSet[id]; !exists {
removed = append(removed, id) removed = append(removed, id)
currentSet[id] = struct{}{}
} }
return true return true
}) })
}
return removed return removed
} }
+69 -20
View File
@@ -36,6 +36,7 @@ type delayedSuccessResponseWriter struct {
mu sync.Mutex mu sync.Mutex
writeCount int writeCount int
body []byte
} }
// deadlineResponseWriter simulates an HTTP/2 response blocked on flow control. // 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() w.mu.Unlock()
if writeCount == 1 { if writeCount == 1 {
// Only the first write is delayed. This simulates a transiently wedged map response: // Only the first write is delayed: once it returns the session drains its
// long enough to make the batcher time out future sends, // backlog (transient) or observes the stop signal (evicted).
// but short enough that the old session can still recover if we leave it alive
w.firstWriteStartedOnce.Do(func() { w.firstWriteStartedOnce.Do(func() {
close(w.firstWriteStarted) 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 return len(data), nil
} }
@@ -160,6 +164,13 @@ func (w *delayedSuccessResponseWriter) WriteCount() int {
return w.writeCount 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 // recordingResponseWriter records the status code and whether anything was
// written, so a test can tell an explicit error response apart from a handler // 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 // 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") "deleting a node must cancel its ephemeral GC timer")
} }
// TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession // TestGitHubIssue3129_BlockedWriteSessionLifecycle
// tests the scenario reported in // tests the scenario reported in
// https://github.com/juanfont/headscale/issues/3129. // https://github.com/juanfont/headscale/issues/3129.
// //
@@ -476,13 +487,25 @@ func TestDeleteNodeCancelsEphemeralGC(t *testing.T) {
// 2. Block the first map write long enough for the session to stop draining // 2. Block the first map write long enough for the session to stop draining
// its buffered map-response channel. // its buffered map-response channel.
// 3. While that write is blocked, queue enough updates to fill the buffered // 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. // channel and make the next batcher send encounter backpressure.
// 4. That stale-send path removes the session from the batcher, so without an // 4. A transient stall must keep the session registered and deliver the queued
// explicit teardown hook the old [mapSession.serveLongPoll] goroutine would stay alive // updates once the write completes. A stall past the batcher's stall budget
// but stop receiving future updates. // must remove the connection and signal its owning session before shutdown.
// 5. Release the blocked write and verify the batcher-side stop signal makes // 5. After the stalled write completes, verify that [mapSession.serveLongPoll]
// that stale session exit instead of lingering as an orphaned goroutine. // exits instead of lingering alive without receiving future updates.
func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t *testing.T) { func TestGitHubIssue3129_BlockedWriteSessionLifecycle(t *testing.T) {
t.Parallel()
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() t.Parallel()
app := createTestApp(t) app := createTestApp(t)
@@ -515,7 +538,7 @@ func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t *
node := nodeView.AsStruct() node := nodeView.AsStruct()
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
writer := newDelayedSuccessResponseWriter(250 * time.Millisecond) writer := newDelayedSuccessResponseWriter(tt.writeDelay)
session := app.newMapSession(ctx, tailcfg.MapRequest{ session := app.newMapSession(ctx, tailcfg.MapRequest{
Stream: true, Stream: true,
Version: tailcfg.CapabilityVersion(100), Version: tailcfg.CapabilityVersion(100),
@@ -551,21 +574,45 @@ func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t *
close(streamsClosed) close(streamsClosed)
}() }()
// One update fills the buffered session channel while the first write is blocked. // Distinct pings cannot collapse into a state refresh on retry. Both must
// The second update then hits the 50ms stale-send timeout, so the batcher prunes // reach a recovering client, including the one that encountered a full channel.
// the stale connection and triggers its stop hook. const (
app.mapBatcher.AddWork(change.SelfUpdate(node.ID), change.SelfUpdate(node.ID)) 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 { select {
case <-writer.FirstWriteFinished(): case <-writer.FirstWriteFinished():
case <-time.After(2 * time.Second): case <-time.After(tt.writeDelay + time.Second):
t.Fatal("expected the blocked write to eventually complete") t.Fatal("expected the blocked write to eventually complete")
} }
require.Eventually(t, session.cancelChClosed.Load, 2*time.Second, 10*time.Millisecond, if !tt.evicted {
"stale eviction must stop the session before batcher shutdown") 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")
// Shutdown skips reconnect grace; the stream wait now includes disconnect cleanup. return
}
// Shutdown skips reconnect grace; the stream wait includes disconnect cleanup.
app.mapBatcher.Close() app.mapBatcher.Close()
assert.Eventually(t, func() bool { assert.Eventually(t, func() bool {
@@ -576,6 +623,8 @@ func TestGitHubIssue3129_TransientlyBlockedWriteDoesNotLeaveLiveStaleSession(t *
return false return false
} }
}, time.Second, 20*time.Millisecond, "after stale-send cleanup, the stale session should exit") }, time.Second, 20*time.Millisecond, "after stale-send cleanup, the stale session should exit")
})
}
} }
// TestDeletedNodeEndsLongPoll proves that deleting a node ends its long-poll // TestDeletedNodeEndsLongPoll proves that deleting a node ends its long-poll