diff --git a/hscontrol/db/ephemeral_garbage_collector_test.go b/hscontrol/db/ephemeral_garbage_collector_test.go index 42472660..a04c2103 100644 --- a/hscontrol/db/ephemeral_garbage_collector_test.go +++ b/hscontrol/db/ephemeral_garbage_collector_test.go @@ -2,6 +2,7 @@ package db import ( "runtime" + "slices" "sync" "sync/atomic" "testing" @@ -9,6 +10,7 @@ import ( "github.com/juanfont/headscale/hscontrol/types" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" ) const ( @@ -119,6 +121,52 @@ func TestEphemeralGarbageCollectorCancelReapsGoroutine(t *testing.T) { }, 2*time.Second, 20*time.Millisecond, "watcher goroutines should be reaped") } +// TestEphemeralGarbageCollectorCancelBeatsQueuedDeletion verifies that a node +// reconnecting (Cancel) after its deletion has already been queued on the +// internal channel is not deleted. The timer fires and enqueues the deletion; +// Cancel then runs before Start drains it. Start must drop the now-superseded +// deletion rather than removing the freshly reconnected node. +func TestEphemeralGarbageCollectorCancelBeatsQueuedDeletion(t *testing.T) { + const targetNode types.NodeID = 42 + + var ( + mu sync.Mutex + deleted []types.NodeID + ) + + e := NewEphemeralGarbageCollector(func(ni types.NodeID) { + mu.Lock() + defer mu.Unlock() + + deleted = append(deleted, ni) + }) + + // Schedule with a tiny expiry but do not drain yet: the watcher fires and + // enqueues the deletion onto the buffered channel. + e.Schedule(targetNode, time.Millisecond) + require.Eventually(t, func() bool { + return len(e.deleteCh) == 1 + }, time.Second, time.Millisecond, "deletion should be queued") + + // Node reconnects before the queue is drained. + e.Cancel(targetNode) + + go e.Start() + defer e.Close() + + require.Eventually(t, func() bool { + return len(e.deleteCh) == 0 + }, time.Second, time.Millisecond, "Start should drain the queued deletion") + + assert.Never(t, func() bool { + mu.Lock() + defer mu.Unlock() + + return slices.Contains(deleted, targetNode) + }, 200*time.Millisecond, 10*time.Millisecond, + "cancelled node must not be deleted") +} + // TestEphemeralGarbageCollectorReschedule is a test for the rescheduling of nodes in [EphemeralGarbageCollector]. // It creates a new [EphemeralGarbageCollector], schedules a node for deletion with a longer expiry, // and then reschedules it with a shorter expiry, and verifies that the node is deleted only once. diff --git a/hscontrol/db/node.go b/hscontrol/db/node.go index fe8d63e9..fd4c8f12 100644 --- a/hscontrol/db/node.go +++ b/hscontrol/db/node.go @@ -394,17 +394,30 @@ type EphemeralGarbageCollector struct { deleteFunc func(types.NodeID) toBeDeleted map[types.NodeID]ephemeralTimer + // gen is bumped for every scheduled deletion so a queued deletion that + // was superseded by a Cancel or reschedule can be recognised and dropped. + gen uint64 - deleteCh chan types.NodeID + deleteCh chan pendingDeletion cancelCh chan struct{} } // ephemeralTimer pairs a node's pending-deletion timer with a done channel -// used to reap its watcher goroutine on Cancel or reschedule. Without it a -// stopped timer never fires and the goroutine leaks until Close. +// used to reap its watcher goroutine on Cancel or reschedule, plus the +// generation identifying this particular scheduling. Without the done channel +// a stopped timer never fires and the goroutine leaks until Close. type ephemeralTimer struct { timer *time.Timer done chan struct{} + gen uint64 +} + +// pendingDeletion is the generation-stamped deletion a watcher enqueues when +// its timer fires. Start drops it if the node's current generation no longer +// matches, i.e. it was cancelled or rescheduled in the meantime. +type pendingDeletion struct { + nodeID types.NodeID + gen uint64 } // NewEphemeralGarbageCollector creates a new [EphemeralGarbageCollector], it takes @@ -412,7 +425,7 @@ type ephemeralTimer struct { func NewEphemeralGarbageCollector(deleteFunc func(types.NodeID)) *EphemeralGarbageCollector { return &EphemeralGarbageCollector{ toBeDeleted: make(map[types.NodeID]ephemeralTimer), - deleteCh: make(chan types.NodeID, 10), + deleteCh: make(chan pendingDeletion, 10), cancelCh: make(chan struct{}), deleteFunc: deleteFunc, } @@ -455,9 +468,11 @@ func (e *EphemeralGarbageCollector) Schedule(nodeID types.NodeID, expiry time.Du close(old.done) } + e.gen++ + gen := e.gen timer := time.NewTimer(expiry) done := make(chan struct{}) - e.toBeDeleted[nodeID] = ephemeralTimer{timer: timer, done: done} + e.toBeDeleted[nodeID] = ephemeralTimer{timer: timer, done: done, gen: gen} // Start a goroutine to handle the timer completion go func() { select { @@ -467,7 +482,7 @@ func (e *EphemeralGarbageCollector) Schedule(nodeID types.NodeID, expiry time.Du // i.e. We don't want to send to deleteCh if the GC is shutting down // So, we try to send to deleteCh, but also watch for cancelCh select { - case e.deleteCh <- nodeID: + case e.deleteCh <- pendingDeletion{nodeID: nodeID, gen: gen}: // Successfully sent to deleteCh case <-e.cancelCh: // GC is shutting down, don't send to deleteCh @@ -504,12 +519,22 @@ func (e *EphemeralGarbageCollector) Start() { select { case <-e.cancelCh: return - case nodeID := <-e.deleteCh: + case pd := <-e.deleteCh: e.mu.Lock() - delete(e.toBeDeleted, nodeID) + + entry, ok := e.toBeDeleted[pd.nodeID] + if !ok || entry.gen != pd.gen { + // Cancelled or rescheduled after this deletion was queued; + // drop it so a reconnected node is not removed. + e.mu.Unlock() + + continue + } + + delete(e.toBeDeleted, pd.nodeID) e.mu.Unlock() - go e.deleteFunc(nodeID) + go e.deleteFunc(pd.nodeID) } } }