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

mafredri · web-flow · commit 41ae01d2e9d9 · 2023-02-14T14:57:48.000Z
diff --git a/agent/apphealth_test.go b/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()) {
diff --git a/coderd/coderdtest/coderdtest.go b/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,
 		})
diff --git a/coderd/workspaceagents.go b/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)
 	}
 
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.DRPCConn().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
`@@ -158,7 +158,7 @@ func TestAppHealth_NotSpamming(t *testing.T) {`
`158`	`158`	`// Ensure we haven't made more than 2 (expected 1 + 1 for buffer) requests in the last second.`
`159`	`159`	`// if there is a bug where we are spamming the healthcheck route this will catch it.`
`160`	`160`	`time.Sleep(time.Second)`
`161`		`- require.LessOrEqual(t, *counter, int32(2))`
	`161`	`+ require.LessOrEqual(t, atomic.LoadInt32(counter), int32(2))`
`162`	`162`	`}`
`163`	`163`
`164`	`164`	`func setupAppReporter(ctx context.Context, t *testing.T, apps []codersdk.WorkspaceApp, handlers []http.Handler) (agent.WorkspaceAgentApps, func()) {`
Original file line number	Diff line number	Diff line change
`@@ -447,6 +447,8 @@ func (api API) dialWorkspaceAgentTailnet(r http.Request, agentID uuid.UUID) (*`
`447`	`447`	`Logger: api.Logger.Named("tailnet"),`
`448`	`448`	`})`
`449`	`449`	`if err != nil {`
	`450`	`+ _ = clientConn.Close()`
	`451`	`+ _ = serverConn.Close()`
`450`	`452`	`return nil, xerrors.Errorf("create tailnet conn: %w", err)`
`451`	`453`	`}`
`452`	`454`
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.DRPCConn().Close()`
	`525`	`+ }`
	`526`	`+`
`523`	`527`	`return err`
`524`	`528`	`}`