fix: Improve closure of provisioner and agent tailnet dial

mafredri · mafredri · commit 64f6496eae5e · 2023-02-14T14:01:28.000Z
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -393,13 +393,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,
 		})
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -446,6 +446,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)
 	}
 
diff --git a/provisionerd/provisionerd.go b/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.Close()
+	}
+
 	return err
 }
diff --git a/provisionerd/provisionerd_test.go b/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{

Original file line number	Diff line number	Diff line change
`@@ -446,6 +446,8 @@ func (api API) dialWorkspaceAgentTailnet(r http.Request, agentID uuid.UUID) (*`
`446`	`446`	`Logger: api.Logger.Named("tailnet"),`
`447`	`447`	`})`
`448`	`448`	`if err != nil {`
	`449`	`+ _ = clientConn.Close()`
	`450`	`+ _ = serverConn.Close()`
`449`	`451`	`return nil, xerrors.Errorf("create tailnet conn: %w", err)`
`450`	`452`	`}`
`451`	`453`
Original file line number	Diff line number	Diff line change
`@@ -520,5 +520,9 @@ func (p *Server) closeWithError(err error) error {`
`520`	`520`
`521`	`521`	`p.opts.Logger.Debug(context.Background(), "closing server with error", slog.Error(err))`
`522`	`522`
	`523`	`+ if c, ok := p.clientValue.Load().(proto.DRPCProvisionerDaemonClient); ok {`
	`524`	`+ _ = c.Close()`
	`525`	`+ }`
	`526`	`+`
`523`	`527`	`return err`
`524`	`528`	`}`