Skip to content

Commit d2e2507

Browse files
johnstcnjohnstcn
committed
address Spike's comments
1 parent 904b857 commit d2e2507

File tree

3 files changed

+24
-39
lines changed

3 files changed

+24
-39
lines changed

‎coderd/coderdtest/coderdtest.go

Lines changed: 6 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -559,8 +559,6 @@ func NewWithAPI(t testing.TB, options *Options) (*codersdk.Client, io.Closer, *c
559559
var provisionerCloser io.Closer = nopcloser{}
560560
if options.IncludeProvisionerDaemon {
561561
provisionerCloser = NewTaggedProvisionerDaemon(t, coderAPI, "test", options.ProvisionerDaemonTags)
562-
// Wait for the provisioner daemon to be ready before continuing.
563-
AwaitProvisionerDaemonsConnected(t, coderAPI)
564562
}
565563
client := codersdk.New(serverURL)
566564
t.Cleanup(func() {
@@ -572,36 +570,6 @@ func NewWithAPI(t testing.TB, options *Options) (*codersdk.Client, io.Closer, *c
572570
return client, provisionerCloser, coderAPI
573571
}
574572

575-
// AwaitProvisionerDaemonsConnected waits for the provisioner daemon to connect.
576-
func AwaitProvisionerDaemonsConnected(t testing.TB, api *coderd.API) {
577-
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
578-
defer cancel()
579-
for {
580-
select {
581-
case <-ctx.Done():
582-
t.Fatal("provisioner daemon did not connect in time")
583-
case <-time.After(testutil.IntervalFast):
584-
// nolint:gocritic // used for testing only
585-
daemons, err := api.Database.GetProvisionerDaemons(dbauthz.AsSystemReadProvisionerDaemons(ctx))
586-
if err != nil {
587-
if errors.Is(err, sql.ErrNoRows) {
588-
t.Logf("no provisioner daemons found yet")
589-
continue
590-
}
591-
require.NoError(t, err)
592-
}
593-
if len(daemons) == 0 {
594-
t.Logf("no provisioner daemons found yet")
595-
continue
596-
}
597-
for _, daemon := range daemons {
598-
t.Logf("found provisioner daemon %q", daemon.Name)
599-
}
600-
return
601-
}
602-
}
603-
}
604-
605573
// ProvisionerdCloser wraps a provisioner daemon as an io.Closer that can be called multiple times
606574
type ProvisionerdCloser struct {
607575
mu sync.Mutex
@@ -663,6 +631,7 @@ func NewTaggedProvisionerDaemon(t testing.TB, coderAPI *coderd.API, name string,
663631
assert.NoError(t, err)
664632
}()
665633

634+
connectedCh := make(chan struct{})
666635
daemon := provisionerd.New(func(dialCtx context.Context) (provisionerdproto.DRPCProvisionerDaemonClient, error) {
667636
return coderAPI.CreateInMemoryTaggedProvisionerDaemon(dialCtx, name, []codersdk.ProvisionerType{codersdk.ProvisionerTypeEcho}, provisionerTags)
668637
}, &provisionerd.Options{
@@ -672,7 +641,12 @@ func NewTaggedProvisionerDaemon(t testing.TB, coderAPI *coderd.API, name string,
672641
Connector: provisionerd.LocalProvisioners{
673642
string(database.ProvisionerTypeEcho): sdkproto.NewDRPCProvisionerClient(echoClient),
674643
},
644+
InitConnectionCh: connectedCh,
675645
})
646+
// Wait for the provisioner daemon to connect before continuing.
647+
// Users of this function tend to assume that the provisioner is connected
648+
// and ready to use when that may not strictly be the case.
649+
<-connectedCh
676650
closer := NewProvisionerDaemonCloser(daemon)
677651
t.Cleanup(func() {
678652
_ = closer.Close()

‎enterprise/coderd/coderdenttest/coderdenttest.go

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -121,8 +121,6 @@ func NewWithAPI(t *testing.T, options *Options) (
121121
var provisionerCloser io.Closer = nopcloser{}
122122
if options.IncludeProvisionerDaemon {
123123
provisionerCloser = coderdtest.NewProvisionerDaemon(t, coderAPI.AGPL)
124-
// Wait for provisioner daemons to be connected before continuing
125-
coderdtest.AwaitProvisionerDaemonsConnected(t, coderAPI.AGPL)
126124
}
127125

128126
t.Cleanup(func() {

‎provisionerd/provisionerd.go

Lines changed: 18 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ type Options struct {
6060
UpdateInterval time.Duration
6161
LogBufferInterval time.Duration
6262
Connector Connector
63+
InitConnectionCh chan struct{} // only to be used in tests
6364
}
6465

6566
// New creates and starts a provisioner daemon.
@@ -84,6 +85,9 @@ func New(clientDialer Dialer, opts *Options) *Server {
8485
mets := NewMetrics(reg)
8586
opts.Metrics = &mets
8687
}
88+
if opts.InitConnectionCh == nil {
89+
opts.InitConnectionCh = make(chan struct{})
90+
}
8791

8892
ctx, ctxCancel := context.WithCancel(context.Background())
8993
daemon := &Server{
@@ -93,11 +97,12 @@ func New(clientDialer Dialer, opts *Options) *Server {
9397
clientDialer: clientDialer,
9498
clientCh: make(chan proto.DRPCProvisionerDaemonClient),
9599

96-
closeContext: ctx,
97-
closeCancel: ctxCancel,
98-
closedCh: make(chan struct{}),
99-
shuttingDownCh: make(chan struct{}),
100-
acquireDoneCh: make(chan struct{}),
100+
closeContext: ctx,
101+
closeCancel: ctxCancel,
102+
closedCh: make(chan struct{}),
103+
shuttingDownCh: make(chan struct{}),
104+
acquireDoneCh: make(chan struct{}),
105+
initConnectionCh: opts.InitConnectionCh,
101106
}
102107

103108
daemon.wg.Add(2)
@@ -115,6 +120,11 @@ type Server struct {
115120

116121
wg sync.WaitGroup
117122

123+
// initConnectionCh will receive when the daemon connects to coderd for the
124+
// first time.
125+
initConnectionCh chan struct{}
126+
initConnectionOnce sync.Once
127+
118128
// mutex protects all subsequent fields
119129
mutex sync.Mutex
120130
// closeContext is canceled when we start closing.
@@ -231,6 +241,9 @@ connectLoop:
231241
}
232242
p.opts.Logger.Info(p.closeContext, "successfully connected to coderd")
233243
retrier.Reset()
244+
p.initConnectionOnce.Do(func() {
245+
close(p.initConnectionCh)
246+
})
234247

235248
// serve the client until we are closed or it disconnects
236249
for {

0 commit comments

Comments
 (0)