fix: Improve closure of provisioner and agent tailnet dial (#6199)

agent/apphealth_test.go

@@ -158,7 +158,7 @@ func TestAppHealth_NotSpamming(t *testing.T) {
 	// Ensure we haven't made more than 2 (expected 1 + 1 for buffer) requests in the last second.
 	// if there is a bug where we are spamming the healthcheck route this will catch it.
 	time.Sleep(time.Second)
-	require.LessOrEqual(t, *counter, int32(2))
+	require.LessOrEqual(t, atomic.LoadInt32(counter), int32(2))
 }
 
 func setupAppReporter(ctx context.Context, t *testing.T, apps []codersdk.WorkspaceApp, handlers []http.Handler) (agent.WorkspaceAgentApps, func()) {

coderd/coderdtest/coderdtest.go

@@ -396,13 +396,16 @@ func NewProvisionerDaemon(t *testing.T, coderAPI *coderd.API) io.Closer {
 func NewExternalProvisionerDaemon(t *testing.T, client *codersdk.Client, org uuid.UUID, tags map[string]string) io.Closer {
 	echoClient, echoServer := provisionersdk.MemTransportPipe()
 	ctx, cancelFunc := context.WithCancel(context.Background())
+	serveDone := make(chan struct{})
 	t.Cleanup(func() {
 		_ = echoClient.Close()
 		_ = echoServer.Close()
 		cancelFunc()
+		<-serveDone
 	})
 	fs := afero.NewMemMapFs()
 	go func() {
+		defer close(serveDone)
 		err := echo.Serve(ctx, fs, &provisionersdk.ServeOptions{
 			Listener: echoServer,
 		})

coderd/workspaceagents.go

@@ -447,6 +447,8 @@ func (api *API) dialWorkspaceAgentTailnet(r *http.Request, agentID uuid.UUID) (*
 		Logger:    api.Logger.Named("tailnet"),
 	})
 	if err != nil {
+		_ = clientConn.Close()
+		_ = serverConn.Close()
 		return nil, xerrors.Errorf("create tailnet conn: %w", err)
 	}
 

provisionerd/provisionerd.go

@@ -520,5 +520,9 @@ func (p *Server) closeWithError(err error) error {
 
 	p.opts.Logger.Debug(context.Background(), "closing server with error", slog.Error(err))
 
+	if c, ok := p.clientValue.Load().(proto.DRPCProvisionerDaemonClient); ok {
+		_ = c.DRPCConn().Close()
+	}
+
 	return err
 }

provisionerd/provisionerd_test.go

@@ -1097,7 +1097,6 @@ func createProvisionerDaemonClient(t *testing.T, server provisionerDaemonTestSer
 	}()
 	t.Cleanup(func() {
 		cancelFunc()
-		_ = serverPipe.Close()
 		<-closed
 	})
 	return proto.NewDRPCProvisionerDaemonClient(clientPipe)
@@ -1117,10 +1116,15 @@ func createProvisionerClient(t *testing.T, server provisionerTestServer) sdkprot
 	require.NoError(t, err)
 	srv := drpcserver.New(mux)
 	ctx, cancelFunc := context.WithCancel(context.Background())
-	t.Cleanup(cancelFunc)
+	closed := make(chan struct{})
 	go func() {
+		defer close(closed)
 		_ = srv.Serve(ctx, serverPipe)
 	}()
+	t.Cleanup(func() {
+		cancelFunc()
+		<-closed
+	})
 	return sdkproto.NewDRPCProvisionerClient(clientPipe)
 }
 
@@ -1150,6 +1154,7 @@ type provisionerDaemonTestServer struct {
 func (p *provisionerDaemonTestServer) AcquireJob(ctx context.Context, empty *proto.Empty) (*proto.AcquiredJob, error) {
 	return p.acquireJob(ctx, empty)
 }
+
 func (p *provisionerDaemonTestServer) CommitQuota(ctx context.Context, com *proto.CommitQuotaRequest) (*proto.CommitQuotaResponse, error) {
 	if p.commitQuota == nil {
 		return &proto.CommitQuotaResponse{