diff --git a/agent/apphealth_test.go b/agent/apphealth_test.go index e09b1ca0cc9cc..20c0d152760fc 100644 --- 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 index 3938c64fd5a85..de14693313e41 100644 --- 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 index f04403070ab0b..e8ab2d4c771ab 100644 --- 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 index a7a1e25cdde43..583826c54c194 100644 --- 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 index fa42f66a25212..ca2899dd3fc1a 100644 --- 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{