Skip to content

Commit 9319c39

Browse files
coadlercoadler
authored
fix: additional provisionerd test double closes (#1267)
1 parent 55ad97b commit 9319c39

File tree

1 file changed

+54
-45
lines changed

1 file changed

+54
-45
lines changed

‎provisionerd/provisionerd_test.go

Lines changed: 54 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -130,8 +130,11 @@ func TestProvisionerd(t *testing.T) {
130130
// Ensures tars with "../../../etc/passwd" as the path
131131
// are not allowed to run, and will fail the job.
132132
t.Parallel()
133-
var complete sync.Once
134-
completeChan := make(chan struct{})
133+
var (
134+
completeChan = make(chan struct{})
135+
completeOnce sync.Once
136+
)
137+
135138
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
136139
return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
137140
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
@@ -150,9 +153,7 @@ func TestProvisionerd(t *testing.T) {
150153
},
151154
updateJob: noopUpdateJob,
152155
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
153-
complete.Do(func() {
154-
close(completeChan)
155-
})
156+
completeOnce.Do(func() { close(completeChan) })
156157
return &proto.Empty{}, nil
157158
},
158159
}), nil
@@ -165,8 +166,11 @@ func TestProvisionerd(t *testing.T) {
165166

166167
t.Run("RunningPeriodicUpdate", func(t *testing.T) {
167168
t.Parallel()
168-
var complete sync.Once
169-
completeChan := make(chan struct{})
169+
var (
170+
completeChan = make(chan struct{})
171+
completeOnce sync.Once
172+
)
173+
170174
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
171175
return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
172176
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
@@ -184,9 +188,7 @@ func TestProvisionerd(t *testing.T) {
184188
}, nil
185189
},
186190
updateJob: func(ctx context.Context, update *proto.UpdateJobRequest) (*proto.UpdateJobResponse, error) {
187-
complete.Do(func() {
188-
close(completeChan)
189-
})
191+
completeOnce.Do(func() { close(completeChan) })
190192
return &proto.UpdateJobResponse{}, nil
191193
},
192194
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
@@ -212,19 +214,18 @@ func TestProvisionerd(t *testing.T) {
212214
didLog atomic.Bool
213215
didAcquireJob atomic.Bool
214216
didDryRun atomic.Bool
217+
completeChan = make(chan struct{})
218+
completeOnce sync.Once
215219
)
216-
var complete sync.Once
217-
completeChan := make(chan struct{})
220+
218221
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
219222
return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
220223
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
221-
if didAcquireJob.Load() {
222-
complete.Do(func() {
223-
close(completeChan)
224-
})
224+
if !didAcquireJob.CAS(false, true) {
225+
completeOnce.Do(func() { close(completeChan) })
225226
return &proto.AcquiredJob{}, nil
226227
}
227-
didAcquireJob.Store(true)
228+
228229
return &proto.AcquiredJob{
229230
JobId: "test",
230231
Provisioner: "someprovisioner",
@@ -315,19 +316,18 @@ func TestProvisionerd(t *testing.T) {
315316
didComplete atomic.Bool
316317
didLog atomic.Bool
317318
didAcquireJob atomic.Bool
319+
completeChan = make(chan struct{})
320+
completeOnce sync.Once
318321
)
319-
var complete sync.Once
320-
completeChan := make(chan struct{})
322+
321323
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
322324
return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
323325
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
324-
if didAcquireJob.Load() {
325-
complete.Do(func() {
326-
close(completeChan)
327-
})
326+
if !didAcquireJob.CAS(false, true) {
327+
completeOnce.Do(func() { close(completeChan) })
328328
return &proto.AcquiredJob{}, nil
329329
}
330-
didAcquireJob.Store(true)
330+
331331
return &proto.AcquiredJob{
332332
JobId: "test",
333333
Provisioner: "someprovisioner",
@@ -386,16 +386,18 @@ func TestProvisionerd(t *testing.T) {
386386
var (
387387
didFail atomic.Bool
388388
didAcquireJob atomic.Bool
389+
completeChan = make(chan struct{})
390+
completeOnce sync.Once
389391
)
390-
completeChan := make(chan struct{})
392+
391393
closer := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
392394
return createProvisionerDaemonClient(t, provisionerDaemonTestServer{
393395
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
394-
if didAcquireJob.Load() {
395-
close(completeChan)
396+
if !didAcquireJob.CAS(false, true) {
397+
completeOnce.Do(func() { close(completeChan) })
396398
return &proto.AcquiredJob{}, nil
397399
}
398-
didAcquireJob.Store(true)
400+
399401
return &proto.AcquiredJob{
400402
JobId: "test",
401403
Provisioner: "someprovisioner",
@@ -585,10 +587,15 @@ func TestProvisionerd(t *testing.T) {
585587

586588
t.Run("ReconnectAndFail", func(t *testing.T) {
587589
t.Parallel()
588-
var second atomic.Bool
589-
failChan := make(chan struct{})
590-
failedChan := make(chan struct{})
591-
completeChan := make(chan struct{})
590+
var (
591+
second atomic.Bool
592+
failChan = make(chan struct{})
593+
failOnce sync.Once
594+
failedChan = make(chan struct{})
595+
failedOnce sync.Once
596+
completeChan = make(chan struct{})
597+
completeOnce sync.Once
598+
)
592599
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
593600
client := createProvisionerDaemonClient(t, provisionerDaemonTestServer{
594601
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
@@ -613,10 +620,10 @@ func TestProvisionerd(t *testing.T) {
613620
},
614621
failJob: func(ctx context.Context, job *proto.FailedJob) (*proto.Empty, error) {
615622
if second.Load() {
616-
close(completeChan)
623+
completeOnce.Do(func() { close(completeChan) })
617624
return &proto.Empty{}, nil
618625
}
619-
close(failChan)
626+
failOnce.Do(func() { close(failChan) })
620627
<-failedChan
621628
return &proto.Empty{}, nil
622629
},
@@ -626,7 +633,7 @@ func TestProvisionerd(t *testing.T) {
626633
<-failChan
627634
_ = client.DRPCConn().Close()
628635
second.Store(true)
629-
close(failedChan)
636+
failedOnce.Do(func() { close(failedChan) })
630637
}()
631638
}
632639
return client, nil
@@ -651,18 +658,20 @@ func TestProvisionerd(t *testing.T) {
651658

652659
t.Run("ReconnectAndComplete", func(t *testing.T) {
653660
t.Parallel()
654-
var completed sync.Once
655-
var second atomic.Bool
656-
failChan := make(chan struct{})
657-
failedChan := make(chan struct{})
658-
completeChan := make(chan struct{})
661+
var (
662+
second atomic.Bool
663+
failChan = make(chan struct{})
664+
failOnce sync.Once
665+
failedChan = make(chan struct{})
666+
failedOnce sync.Once
667+
completeChan = make(chan struct{})
668+
completeOnce sync.Once
669+
)
659670
server := createProvisionerd(t, func(ctx context.Context) (proto.DRPCProvisionerDaemonClient, error) {
660671
client := createProvisionerDaemonClient(t, provisionerDaemonTestServer{
661672
acquireJob: func(ctx context.Context, _ *proto.Empty) (*proto.AcquiredJob, error) {
662673
if second.Load() {
663-
completed.Do(func() {
664-
close(completeChan)
665-
})
674+
completeOnce.Do(func() { close(completeChan) })
666675
return &proto.AcquiredJob{}, nil
667676
}
668677
return &proto.AcquiredJob{
@@ -688,7 +697,7 @@ func TestProvisionerd(t *testing.T) {
688697
if second.Load() {
689698
return &proto.Empty{}, nil
690699
}
691-
close(failChan)
700+
failOnce.Do(func() { close(failChan) })
692701
<-failedChan
693702
return &proto.Empty{}, nil
694703
},
@@ -698,7 +707,7 @@ func TestProvisionerd(t *testing.T) {
698707
<-failChan
699708
_ = client.DRPCConn().Close()
700709
second.Store(true)
701-
close(failedChan)
710+
failedOnce.Do(func() { close(failedChan) })
702711
}()
703712
}
704713
return client, nil

0 commit comments

Comments
 (0)