fix: Remove race condition with acquiredJobDone channel (#148)

bryphe-coder · web-flow · commit 9cc4756ef68b · 2022-02-02T17:32:20.000-08:00
Found another data race while running the tests: https://github.com/coder/coder/runs/5044320845?check_suite_focus=true#step:7:83 __Issue:__ There is a race in the p.acquiredJobDone chan - in particular, there can be a case where we're waiting on the channel to finish (in close) with <-p.acquiredJobDone, but in parallel, an acquireJob could've been started, which would create a new channel for p.acquiredJobDone. There is a similar race in `close(..)`ing the channel, which also came up in test runs. __Fix:__ Instead of recreating the channel everytime, we can use `sync.WaitGroup` to accomplish the same functionality - a semaphore to make close wait for the current job to wrap up.
diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go
@@ -90,7 +90,7 @@ type provisionerDaemon struct {
 	acquiredJobCancel    context.CancelFunc
 	acquiredJobCancelled atomic.Bool
 	acquiredJobRunning   atomic.Bool
-	acquiredJobDone      chan struct{}
+	acquiredJobGroup     sync.WaitGroup
 }
 
 // Connect establishes a connection to coderd.
@@ -184,7 +184,7 @@ func (p *provisionerDaemon) acquireJob(ctx context.Context) {
 	ctx, p.acquiredJobCancel = context.WithCancel(ctx)
 	p.acquiredJobCancelled.Store(false)
 	p.acquiredJobRunning.Store(true)
-	p.acquiredJobDone = make(chan struct{})
+	p.acquiredJobGroup.Add(1)
 
 	p.opts.Logger.Info(context.Background(), "acquired job",
 		slog.F("organization_name", p.acquiredJob.OrganizationName),
@@ -235,7 +235,7 @@ func (p *provisionerDaemon) runJob(ctx context.Context) {
 		p.acquiredJobMutex.Lock()
 		defer p.acquiredJobMutex.Unlock()
 		p.acquiredJobRunning.Store(false)
-		close(p.acquiredJobDone)
+		p.acquiredJobGroup.Done()
 	}()
 	// It's safe to cast this ProvisionerType. This data is coming directly from coderd.
 	provisioner, hasProvisioner := p.opts.Provisioners[p.acquiredJob.Provisioner]
@@ -520,7 +520,7 @@ func (p *provisionerDaemon) closeWithError(err error) error {
 		if !p.acquiredJobCancelled.Load() {
 			p.cancelActiveJob(errMsg)
 		}
-		<-p.acquiredJobDone
+		p.acquiredJobGroup.Wait()
 	}
 
 	p.opts.Logger.Debug(context.Background(), "closing server with error", slog.Error(err))
