fix: track goroutines and fix race condition in reconciler (#21980)

## Problem

CI failure showed 3 goroutines leaked in the prebuilds reconciler, all
stuck in `select` state:

1) `MetricsCollector.BackgroundFetch` (metrics goroutine)
2) `StoreReconciler.Run` (main reconciliation loop)
3) `StoreReconciler.Run.func3()` (provisioner job publisher goroutine)

All three goroutines were waiting for `ctx.Done()`, which likely means
`cancelFn()` was never called to trigger shutdown.

**Note:** I was unable to reproduce the flake locally. The likely cause
was a race condition between `Run()` and `Stop()` where `Stop()` could
check `running` (seeing `false`), return early, and then `Run()` would
start goroutines that never get cleaned up. This could happen in any
`coderd` test that starts a server with prebuilds enabled.

### Problems identified

1) Missing waitgoroup tracking: provisioner job publisher goroutine was
not tracked in the waitgroup, therefore, this goroutine was not tracked
for a clean shutdown in `Run defer func()`.
2) The provisioner job publisher goroutine had a redundant `case
<-c.done` that could race with `Stop()` select statement.
3) Race condition between `Run()` and `Stop()`: the `running` and
`stopped` fields were `atomic.Bool` values checked and set
independently, allowing a window where `Stop()` could see
`running=false` and return early, then `Run()` would set `running=true`
and start goroutines that would never be cleaned up. This could happen
in any `coderd` test that starts a server with prebuilds enabled.

## Changes

* Added `wg.Add(1)` and `defer wg.Done()` to track provisioner job
publisher goroutine in waitgroup
* Removed redundant `case <-c.done` from provisioner job publisher
goroutine to eliminate race condition
* Replaced `atomic.Bool` for `running` and `stopped` with a `sync.Mutex`
lifecycle state, also protecting `cancelFn` under the same mutex, to
eliminate the race between `Run()` and `Stop()`
* Added a guard in `Run()` to prevent double-start (`c.stopped ||
c.running`)
* Improved comments in Stop() and Run() to clarify shutdown behavior

Closes: https://github.com/coder/internal/issues/1116

enterprise/coderd/prebuilds/reconcile.go

@@ -51,9 +51,12 @@ type StoreReconciler struct {
 	buildUsageChecker *atomic.Pointer[wsbuilder.UsageChecker]
 	tracer            trace.Tracer
 
-	cancelFn          context.CancelCauseFunc
-	running           atomic.Bool
-	stopped           atomic.Bool
+	// mu protects the reconciler's lifecycle state.
+	mu       sync.Mutex
+	running  bool
+	stopped  bool
+	cancelFn context.CancelCauseFunc
+
 	done              chan struct{}
 	provisionNotifyCh chan database.ProvisionerJob
 
@@ -174,18 +177,33 @@ func (c *StoreReconciler) Run(ctx context.Context) {
 		slog.F("backoff_lookback", c.cfg.ReconciliationBackoffLookback.String()),
 		slog.F("preset_concurrency", c.reconciliationConcurrency))
 
-	var wg sync.WaitGroup
+	// Create a child context that will be canceled when:
+	// 1. The parent context is canceled, OR
+	// 2. c.cancelFn() is called to trigger shutdown
+	// nolint:gocritic // Reconciliation Loop needs Prebuilds Orchestrator permissions.
+	ctx, cancel := context.WithCancelCause(dbauthz.AsPrebuildsOrchestrator(ctx))
+
+	// If the reconciler was already stopped, exit early and release the context.
+	// Otherwise, mark it as running and store the cancel function for shutdown.
+	c.mu.Lock()
+	if c.stopped || c.running {
+		c.mu.Unlock()
+		cancel(nil)
+		return
+	}
+	c.running = true
+	c.cancelFn = cancel
+	c.mu.Unlock()
+
 	ticker := c.clock.NewTicker(reconciliationInterval)
 	defer ticker.Stop()
+	// Wait for all background goroutines to exit before signaling completion.
+	var wg sync.WaitGroup
 	defer func() {
 		wg.Wait()
 		c.done <- struct{}{}
 	}()
 
-	// nolint:gocritic // Reconciliation Loop needs Prebuilds Orchestrator permissions.
-	ctx, cancel := context.WithCancelCause(dbauthz.AsPrebuildsOrchestrator(ctx))
-	c.cancelFn = cancel
-
 	// Start updating metrics in the background.
 	if c.metrics != nil {
 		wg.Add(1)
@@ -195,11 +213,6 @@ func (c *StoreReconciler) Run(ctx context.Context) {
 		}()
 	}
 
-	// Everything is in place, reconciler can now be considered as running.
-	//
-	// NOTE: without this atomic bool, Stop might race with Run for the c.cancelFn above.
-	c.running.Store(true)
-
 	// Publish provisioning jobs outside of database transactions.
 	// A connection is held while a database transaction is active; PGPubsub also tries to acquire a new connection on
 	// Publish, so we can exhaust available connections.
@@ -207,11 +220,11 @@ func (c *StoreReconciler) Run(ctx context.Context) {
 	// A single worker dequeues from the channel, which should be sufficient.
 	// If any messages are missed due to congestion or errors, provisionerdserver has a backup polling mechanism which
 	// will periodically pick up any queued jobs (see poll(time.Duration) in coderd/provisionerdserver/acquirer.go).
+	wg.Add(1)
 	go func() {
+		defer wg.Done()
 		for {
 			select {
-			case <-c.done:
-				return
 			case <-ctx.Done():
 				return
 			case job := <-c.provisionNotifyCh:
@@ -256,21 +269,29 @@ func (c *StoreReconciler) Run(ctx context.Context) {
 	}
 }
 
+// Stop triggers reconciler shutdown and waits for it to complete.
+// The ctx parameter provides a timeout, if cleanup doesn't finish within
+// this timeout, Stop() logs an error and returns.
 func (c *StoreReconciler) Stop(ctx context.Context, cause error) {
-	defer c.running.Store(false)
-
 	if cause != nil {
 		c.logger.Info(context.Background(), "stopping reconciler", slog.F("cause", cause.Error()))
 	} else {
 		c.logger.Info(context.Background(), "stopping reconciler")
 	}
 
-	// If previously stopped (Swap returns previous value), then short-circuit.
+	// Mark the reconciler as stopped. If it was already stopped, return early.
+	// If the reconciler is running, we'll proceed to shut it down.
 	//
-	// NOTE: we need to *prospectively* mark this as stopped to prevent Stop being called multiple times and causing problems.
-	if c.stopped.Swap(true) {
+	// NOTE: we need to *prospectively* mark this as stopped to prevent the
+	// reconciler from being stopped multiple times and causing problems.
+	c.mu.Lock()
+	if c.stopped {
+		c.mu.Unlock()
 		return
 	}
+	c.stopped = true
+	running := c.running
+	c.mu.Unlock()
 
 	// Unregister prebuilds state and operational metrics.
 	if c.metrics != nil && c.registerer != nil {
@@ -289,16 +310,18 @@ func (c *StoreReconciler) Stop(ctx context.Context, cause error) {
 	}
 
 	// If the reconciler is not running, there's nothing else to do.
-	if !c.running.Load() {
+	if !running {
 		return
 	}
 
+	// Trigger reconciler shutdown by canceling its internal context.
 	if c.cancelFn != nil {
 		c.cancelFn(cause)
 	}
 
+	// Wait for the reconciler to signal that it has fully exited and cleaned up.
 	select {
-	// Give up waiting for control loop to exit.
+	// Timeout: reconciler didn't finish cleanup within the timeout period.
 	case <-ctx.Done():
 		// nolint:gocritic // it's okay to use slog.F() for an error in this case
 		// because we want to differentiate two different types of errors: ctx.Err() and context.Cause()
@@ -308,7 +331,7 @@ func (c *StoreReconciler) Stop(ctx context.Context, cause error) {
 			slog.Error(ctx.Err()),
 			slog.F("cause", context.Cause(ctx)),
 		)
-	// Wait for the control loop to exit.
+	// Happy path: reconciler has successfully exited.
 	case <-c.done:
 		c.logger.Info(context.Background(), "reconciler stopped")
 	}

enterprise/coderd/prebuilds/reconcile_test.go

@@ -1279,9 +1279,8 @@ func TestRunLoop(t *testing.T) {
 		ReconciliationBackoffInterval: serpent.Duration(backoffInterval),
 		ReconciliationInterval:        serpent.Duration(time.Second),
 	}
-	logger := slogtest.Make(
-		t, &slogtest.Options{IgnoreErrors: true},
-	).Leveled(slog.LevelDebug)
+	// Do not ignore errors as we want a graceful stop
+	logger := slogtest.Make(t, &slogtest.Options{IgnoreErrors: false}).Leveled(slog.LevelDebug)
 	db, pubSub := dbtestutil.NewDB(t)
 	cache := files.New(prometheus.NewRegistry(), &coderdtest.FakeAuthorizer{})
 	reconciler := prebuilds.NewStoreReconciler(