Protect race on cancelFn by using running/stopped atomic bools

dannykopping · dannykopping · commit 9676b658d67b · 2025-04-24T16:13:51.000+02:00
Signed-off-by: Danny Kopping &lt;dannykopping@gmail.com&gt;
diff --git a/coderd/prebuilds/api.go b/coderd/prebuilds/api.go
@@ -14,10 +14,10 @@ var ErrNoClaimablePrebuiltWorkspaces = xerrors.New("no claimable prebuilt worksp
 type ReconciliationOrchestrator interface {
 	Reconciler
 
-	// RunLoop starts a continuous reconciliation loop that periodically calls ReconcileAll
+	// Run starts a continuous reconciliation loop that periodically calls ReconcileAll
 	// to ensure all prebuilds are in their desired states. The loop runs until the context
 	// is canceled or Stop is called.
-	RunLoop(ctx context.Context)
+	Run(ctx context.Context)
 
 	// Stop gracefully shuts down the orchestrator with the given cause.
 	// The cause is used for logging and error reporting.
diff --git a/coderd/prebuilds/noop.go b/coderd/prebuilds/noop.go
@@ -10,7 +10,7 @@ import (
 
 type NoopReconciler struct{}
 
-func (NoopReconciler) RunLoop(context.Context)            {}
+func (NoopReconciler) Run(context.Context)                {}
 func (NoopReconciler) Stop(context.Context, error)        {}
 func (NoopReconciler) ReconcileAll(context.Context) error { return nil }
 func (NoopReconciler) SnapshotState(context.Context, database.Store) (*GlobalSnapshot, error) {
diff --git a/enterprise/coderd/coderd.go b/enterprise/coderd/coderd.go
@@ -874,7 +874,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {
 			}
 
 			api.AGPL.PrebuildsReconciler.Store(&reconciler)
-			go reconciler.RunLoop(context.Background())
+			go reconciler.Run(context.Background())
 
 			api.AGPL.PrebuildsClaimer.Store(&claimer)
 		}
diff --git a/enterprise/coderd/prebuilds/reconcile.go b/enterprise/coderd/prebuilds/reconcile.go
@@ -38,6 +38,7 @@ type StoreReconciler struct {
 	clock  quartz.Clock
 
 	cancelFn context.CancelCauseFunc
+	running  atomic.Bool
 	stopped  atomic.Bool
 	done     chan struct{}
 }
@@ -61,7 +62,7 @@ func NewStoreReconciler(
 	}
 }
 
-func (c *StoreReconciler) RunLoop(ctx context.Context) {
+func (c *StoreReconciler) Run(ctx context.Context) {
 	reconciliationInterval := c.cfg.ReconciliationInterval.Value()
 	if reconciliationInterval <= 0 { // avoids a panic
 		reconciliationInterval = 5 * time.Minute
@@ -82,6 +83,11 @@ func (c *StoreReconciler) RunLoop(ctx context.Context) {
 	ctx, cancel := context.WithCancelCause(dbauthz.AsPrebuildsOrchestrator(ctx))
 	c.cancelFn = cancel
 
+	// 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)
+
 	for {
 		select {
 		// TODO: implement pubsub listener to allow reconciling a specific template imperatively once it has been changed,
@@ -107,16 +113,26 @@ func (c *StoreReconciler) RunLoop(ctx context.Context) {
 }
 
 func (c *StoreReconciler) Stop(ctx context.Context, cause error) {
+	defer c.running.Store(false)
+
 	if cause != nil {
 		c.logger.Error(context.Background(), "stopping reconciler due to an error", slog.Error(cause))
 	} else {
 		c.logger.Info(context.Background(), "gracefully stopping reconciler")
 	}
 
-	if c.isStopped() {
+	// If previously stopped (Swap returns previous value), then short-circuit.
+	//
+	// NOTE: we need to *prospectively* mark this as stopped to prevent Stop being called multiple times and causing problems.
+	if c.stopped.Swap(true) {
+		return
+	}
+
+	// If the reconciler is not running, there's nothing else to do.
+	if !c.running.Load() {
 		return
 	}
-	c.stopped.Store(true)
+
 	if c.cancelFn != nil {
 		c.cancelFn(cause)
 	}
@@ -138,10 +154,6 @@ func (c *StoreReconciler) Stop(ctx context.Context, cause error) {
 	}
 }
 
-func (c *StoreReconciler) isStopped() bool {
-	return c.stopped.Load()
-}
-
 // ReconcileAll will attempt to resolve the desired vs actual state of all templates which have presets with prebuilds configured.
 //
 // NOTE:
diff --git a/enterprise/coderd/prebuilds/reconcile_test.go b/enterprise/coderd/prebuilds/reconcile_test.go
@@ -575,7 +575,7 @@ func TestRunLoop(t *testing.T) {
 		t, &slogtest.Options{IgnoreErrors: true},
 	).Leveled(slog.LevelDebug)
 	db, pubSub := dbtestutil.NewDB(t)
-	controller := prebuilds.NewStoreReconciler(db, pubSub, cfg, logger, clock)
+	reconciler := prebuilds.NewStoreReconciler(db, pubSub, cfg, logger, clock)
 
 	ownerID := uuid.New()
 	dbgen.User(t, db, database.User{
@@ -639,7 +639,7 @@ func TestRunLoop(t *testing.T) {
 	// we need to wait until ticker is initialized, and only then use clock.Advance()
 	// otherwise clock.Advance() will be ignored
 	trap := clock.Trap().NewTicker()
-	go controller.RunLoop(ctx)
+	go reconciler.Run(ctx)
 	// wait until ticker is initialized
 	trap.MustWait(ctx).Release()
 	// start 1st iteration of ReconciliationLoop
@@ -681,7 +681,7 @@ func TestRunLoop(t *testing.T) {
 	}, testutil.WaitShort, testutil.IntervalFast)
 
 	// gracefully stop the reconciliation loop
-	controller.Stop(ctx, nil)
+	reconciler.Stop(ctx, nil)
 }
 
 func TestFailedBuildBackoff(t *testing.T) {

Original file line number	Diff line number	Diff line change
`@@ -874,7 +874,7 @@ func (api *API) updateEntitlements(ctx context.Context) error {`
`874`	`874`	`}`
`875`	`875`
`876`	`876`	`api.AGPL.PrebuildsReconciler.Store(&reconciler)`
`877`		`- go reconciler.RunLoop(context.Background())`
	`877`	`+ go reconciler.Run(context.Background())`
`878`	`878`
`879`	`879`	`api.AGPL.PrebuildsClaimer.Store(&claimer)`
`880`	`880`	`}`