@@ -50,12 +50,11 @@ func New(clientDialer Dialer, opts *Options) io.Closer {
5050}
5151ctx ,ctxCancel := context .WithCancel (context .Background ())
5252daemon := & provisionerDaemon {
53- clientDialer :clientDialer ,
54- opts :opts ,
55- acquiredJobDone :make (chan struct {}),
56- closeContext :ctx ,
57- closeCancel :ctxCancel ,
58- closed :make (chan struct {}),
53+ clientDialer :clientDialer ,
54+ opts :opts ,
55+ closeContext :ctx ,
56+ closeCancel :ctxCancel ,
57+ closed :make (chan struct {}),
5958}
6059go daemon .connect (ctx )
6160return daemon
@@ -90,7 +89,7 @@ type provisionerDaemon struct {
9089acquiredJobCancel context.CancelFunc
9190acquiredJobCancelled atomic.Bool
9291acquiredJobRunning atomic.Bool
93- acquiredJobDone chan struct {}
92+ acquiredJobGroup sync. WaitGroup
9493}
9594
9695// Connect establishes a connection to coderd.
@@ -182,6 +181,7 @@ func (p *provisionerDaemon) acquireJob(ctx context.Context) {
182181return
183182}
184183ctx ,p .acquiredJobCancel = context .WithCancel (ctx )
184+ p .acquiredJobGroup .Add (1 )
185185p .acquiredJobCancelled .Store (false )
186186p .acquiredJobRunning .Store (true )
187187
@@ -234,7 +234,7 @@ func (p *provisionerDaemon) runJob(ctx context.Context) {
234234}
235235p .opts .Logger .Debug (ctx ,"cleaned up work directory" )
236236p .acquiredJobRunning .Store (false )
237- close ( p . acquiredJobDone )
237+ p . acquiredJobGroup . Done ( )
238238}()
239239// It's safe to cast this ProvisionerType. This data is coming directly from coderd.
240240provisioner ,hasProvisioner := p .opts .Provisioners [p .acquiredJob .Provisioner ]
@@ -520,7 +520,7 @@ func (p *provisionerDaemon) closeWithError(err error) error {
520520if ! p .acquiredJobCancelled .Load () {
521521p .cancelActiveJob (errMsg )
522522}
523- <- p . acquiredJobDone
523+ p . acquiredJobGroup . Wait ()
524524}
525525
526526p .opts .Logger .Debug (context .Background (),"closing server with error" ,slog .Error (err ))