diff --git a/enterprise/coderd/prebuilds/reconcile.go b/enterprise/coderd/prebuilds/reconcile.go index 6816ce1799..b41b5541fc 100644 --- a/enterprise/coderd/prebuilds/reconcile.go +++ b/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") } diff --git a/enterprise/coderd/prebuilds/reconcile_test.go b/enterprise/coderd/prebuilds/reconcile_test.go index 393322cb1b..d0d40666ac 100644 --- a/enterprise/coderd/prebuilds/reconcile_test.go +++ b/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(