db: drop superseded ephemeral deletions in the GC drain loop

A node reconnecting after its deletion queued is no longer removed.

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.

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)
 		}
 	}
 }