chore(coderd/coderdtest): wait for provisioner daemons to be connected (#15936)

johnstcn · web-flow · commit 9e9a5fd2043b · 2024-12-20T09:49:19.000Z
Fixes coder/internal#260
diff --git 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
@@ -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
@@ -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 {
