wgengine,util/execqueue: wait for in-flight linkChange before closing

ExecQueue.Shutdown does not wait for a function that is already
executing, so Close could tear down magicConn, dns, wgdev, and tundev
while a queued linkChange was still using them, panicking during
shutdown. Add ExecQueue.ShutdownAndWait, which discards queued
functions that have not started and waits for the in-flight one, and
use it in Close with a bounded context before tearing anything down.
The eventbus client is closed first and is the queue's only producer,
so no new work can arrive after the drain.

Updates #17641

Change-Id: I0350bcb59c1ee4b0dcac88cf66b93828466c8c98
Signed-off-by: Adel-Ayoub <adelayoub.maaziz@gmail.com>

util/execqueue/execqueue.go

@@ -88,6 +88,36 @@ func (q *ExecQueue) Shutdown() {
 	}
 }
 
+// ShutdownAndWait signals the queue to stop, discards any queued
+// functions that have not started, and waits for the currently
+// executing function, if any, to complete or ctx to expire.
+//
+// It must not be called while holding a lock that a queued function
+// may acquire, or it will not return until ctx expires.
+func (q *ExecQueue) ShutdownAndWait(ctx context.Context) error {
+	q.mu.Lock()
+	q.closed = true
+	if q.cancel != nil {
+		q.cancel()
+	}
+	waitCh := q.doneWaiter
+	if q.inFlight && waitCh == nil {
+		waitCh = make(chan struct{})
+		q.doneWaiter = waitCh
+	}
+	q.mu.Unlock()
+
+	if waitCh == nil {
+		return nil
+	}
+	select {
+	case <-waitCh:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 func (q *ExecQueue) initCtxLocked() {
 	if q.ctx == nil {
 		q.ctx, q.cancel = context.WithCancel(context.Background())

util/execqueue/execqueue_test.go

@@ -7,6 +7,8 @@ import (
 	"context"
 	"sync/atomic"
 	"testing"
+	"testing/synctest"
+	"time"
 )
 
 func TestExecQueue(t *testing.T) {
@@ -29,3 +31,55 @@ func TestExecQueueRunSyncLocking(t *testing.T) {
 		q.Shutdown()
 	})
 }
+
+func TestShutdownAndWait(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		q := &ExecQueue{}
+		started := make(chan struct{})
+		release := make(chan struct{})
+		var finished, ranPending atomic.Bool
+		q.Add(func() {
+			close(started)
+			<-release
+			finished.Store(true)
+		})
+		q.Add(func() { ranPending.Store(true) })
+		<-started
+
+		// The fake clock only advances once ShutdownAndWait below is
+		// blocked, so the release cannot fire early.
+		go func() {
+			time.Sleep(time.Second)
+			close(release)
+		}()
+		if err := q.ShutdownAndWait(context.Background()); err != nil {
+			t.Fatalf("ShutdownAndWait: %v", err)
+		}
+		if !finished.Load() {
+			t.Error("ShutdownAndWait returned before the in-flight function completed")
+		}
+		if ranPending.Load() {
+			t.Error("pending function ran after shutdown")
+		}
+	})
+}
+
+func TestShutdownAndWaitTimeout(t *testing.T) {
+	synctest.Test(t, func(t *testing.T) {
+		q := &ExecQueue{}
+		started := make(chan struct{})
+		release := make(chan struct{})
+		q.Add(func() {
+			close(started)
+			<-release
+		})
+		<-started
+
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		if err := q.ShutdownAndWait(ctx); err == nil {
+			t.Error("ShutdownAndWait = nil; want deadline exceeded")
+		}
+		close(release)
+	})
+}

wgengine/userspace.go

@@ -1222,9 +1222,14 @@ func (e *userspaceEngine) RequestStatus() {
 
 func (e *userspaceEngine) Close() {
 	e.eventClient.Close()
-	// TODO(cmol): Should we wait for it too?
+	// eventClient.Close waited for the ChangeDelta subscriber, the sole
-	// Same question raised in appconnector.go.
+	// producer for linkChangeQueue, to return, so no new work can be
-	e.linkChangeQueue.Shutdown()
+	// queued. Discard queued linkChanges and wait for an in-flight one
+	// to finish before closing the subsystems it uses.
+	// See tailscale/tailscale#17641.
+	drainCtx, drainCancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer drainCancel()
+	e.linkChangeQueue.ShutdownAndWait(drainCtx)
 	e.mu.Lock()
 	if e.closing {
 		e.mu.Unlock()

wgengine/userspace_test.go

@@ -12,6 +12,7 @@ import (
 	"slices"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/tailscale/wireguard-go/device"
 	"go4.org/mem"
@@ -643,3 +644,40 @@ func TestLinkChangeReapplyPreservesMagicDNSRoutes(t *testing.T) {
 			logger.AsJSON(initial), logger.AsJSON(after))
 	}
 }
+
+// TestCloseWaitsForLinkChange tests that Close waits for in-flight
+// linkChangeQueue work to finish before tearing down the subsystems
+// that linkChange uses.
+//
+// See https://github.com/tailscale/tailscale/issues/17641.
+func TestCloseWaitsForLinkChange(t *testing.T) {
+	bus := eventbustest.NewBus(t)
+
+	ht := health.NewTracker(bus)
+	reg := new(usermetric.Registry)
+	e, err := NewFakeUserspaceEngine(t.Logf, 0, ht, reg, bus)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	started := make(chan struct{})
+	release := make(chan struct{})
+	done := make(chan struct{})
+	e.(*userspaceEngine).linkChangeQueue.Add(func() {
+		close(started)
+		<-release
+		close(done)
+	})
+	<-started
+
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		close(release)
+	}()
+	e.Close()
+	select {
+	case <-done:
+	default:
+		t.Fatal("Close returned with link change work still in flight")
+	}
+}