diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go index 7c1e6a4962a8c..bd1ed740a7ce7 100644 --- a/coderd/coderdtest/coderdtest.go +++ b/coderd/coderdtest/coderdtest.go @@ -631,6 +631,7 @@ func NewTaggedProvisionerDaemon(t testing.TB, coderAPI *coderd.API, name string, assert.NoError(t, err) }() + connectedCh := make(chan struct{}) daemon := provisionerd.New(func(dialCtx context.Context) (provisionerdproto.DRPCProvisionerDaemonClient, error) { return coderAPI.CreateInMemoryTaggedProvisionerDaemon(dialCtx, name, []codersdk.ProvisionerType{codersdk.ProvisionerTypeEcho}, provisionerTags) }, &provisionerd.Options{ @@ -640,7 +641,12 @@ func NewTaggedProvisionerDaemon(t testing.TB, coderAPI *coderd.API, name string, Connector: provisionerd.LocalProvisioners{ string(database.ProvisionerTypeEcho): sdkproto.NewDRPCProvisionerClient(echoClient), }, + InitConnectionCh: connectedCh, }) + // Wait for the provisioner daemon to connect before continuing. + // Users of this function tend to assume that the provisioner is connected + // and ready to use when that may not strictly be the case. + <-connectedCh closer := NewProvisionerDaemonCloser(daemon) t.Cleanup(func() { _ = closer.Close() diff --git a/coderd/templateversions_test.go b/coderd/templateversions_test.go index d8377821245bf..1a67508880188 100644 --- a/coderd/templateversions_test.go +++ b/coderd/templateversions_test.go @@ -172,8 +172,8 @@ func TestPostTemplateVersionsByOrganization(t *testing.T) { require.Equal(t, provisionersdk.ScopeOrganization, version.Job.Tags[provisionersdk.TagScope]) if assert.Equal(t, version.Job.Status, codersdk.ProvisionerJobPending) { assert.NotNil(t, version.MatchedProvisioners) - assert.Equal(t, version.MatchedProvisioners.Available, 1) - assert.Equal(t, version.MatchedProvisioners.Count, 1) + assert.Equal(t, 1, version.MatchedProvisioners.Available) + assert.Equal(t, 1, version.MatchedProvisioners.Count) assert.True(t, version.MatchedProvisioners.MostRecentlySeen.Valid) } diff --git a/provisionerd/provisionerd.go b/provisionerd/provisionerd.go index 4cf2436ea1267..e3b8da8bfe2d9 100644 --- a/provisionerd/provisionerd.go +++ b/provisionerd/provisionerd.go @@ -60,6 +60,7 @@ type Options struct { UpdateInterval time.Duration LogBufferInterval time.Duration Connector Connector + InitConnectionCh chan struct{} // only to be used in tests } // New creates and starts a provisioner daemon. @@ -84,6 +85,9 @@ func New(clientDialer Dialer, opts *Options) *Server { mets := NewMetrics(reg) opts.Metrics = &mets } + if opts.InitConnectionCh == nil { + opts.InitConnectionCh = make(chan struct{}) + } ctx, ctxCancel := context.WithCancel(context.Background()) daemon := &Server{ @@ -93,11 +97,12 @@ func New(clientDialer Dialer, opts *Options) *Server { clientDialer: clientDialer, clientCh: make(chan proto.DRPCProvisionerDaemonClient), - closeContext: ctx, - closeCancel: ctxCancel, - closedCh: make(chan struct{}), - shuttingDownCh: make(chan struct{}), - acquireDoneCh: make(chan struct{}), + closeContext: ctx, + closeCancel: ctxCancel, + closedCh: make(chan struct{}), + shuttingDownCh: make(chan struct{}), + acquireDoneCh: make(chan struct{}), + initConnectionCh: opts.InitConnectionCh, } daemon.wg.Add(2) @@ -115,6 +120,11 @@ type Server struct { wg sync.WaitGroup + // initConnectionCh will receive when the daemon connects to coderd for the + // first time. + initConnectionCh chan struct{} + initConnectionOnce sync.Once + // mutex protects all subsequent fields mutex sync.Mutex // closeContext is canceled when we start closing. @@ -231,6 +241,9 @@ connectLoop: } p.opts.Logger.Info(p.closeContext, "successfully connected to coderd") retrier.Reset() + p.initConnectionOnce.Do(func() { + close(p.initConnectionCh) + }) // serve the client until we are closed or it disconnects for {