From e7d978e0b630969e58b2650f8641abd9ca883f9c Mon Sep 17 00:00:00 2001
From: Colin Adler <colin1adler@gmail.com>
Date: Wed, 9 Aug 2023 23:38:47 +0000
Subject: [PATCH] fix: only init `clistat.Client` when calling `coder stat`

---
 cli/stat.go | 110 ++++++++++++++++++++++++++++++++--------------------
 1 file changed, 67 insertions(+), 43 deletions(-)

diff --git a/cli/stat.go b/cli/stat.go
index 6311596ecea2a..6a49ffb8437b5 100644
--- a/cli/stat.go
+++ b/cli/stat.go
@@ -12,31 +12,43 @@ import (
 	"github.com/coder/coder/cli/cliui"
 )
 
-func (r *RootCmd) stat() *clibase.Cmd {
-	fs := afero.NewReadOnlyFs(afero.NewOsFs())
-	defaultCols := []string{
-		"host_cpu",
-		"host_memory",
-		"home_disk",
-		"container_cpu",
-		"container_memory",
-	}
-	formatter := cliui.NewOutputFormatter(
-		cliui.TableFormat([]statsRow{}, defaultCols),
-		cliui.JSONFormat(),
-	)
-	st, err := clistat.New(clistat.WithFS(fs))
-	if err != nil {
-		panic(xerrors.Errorf("initialize workspace stats collector: %w", err))
+func initStatterMW(tgt **clistat.Statter, fs afero.Fs) clibase.MiddlewareFunc {
+	return func(next clibase.HandlerFunc) clibase.HandlerFunc {
+		return func(i *clibase.Invocation) error {
+			var err error
+			stat, err := clistat.New(clistat.WithFS(fs))
+			if err != nil {
+				return xerrors.Errorf("initialize workspace stats collector: %w", err)
+			}
+			*tgt = stat
+			return next(i)
+		}
 	}
+}
 
+func (r *RootCmd) stat() *clibase.Cmd {
+	var (
+		st        *clistat.Statter
+		fs        = afero.NewReadOnlyFs(afero.NewOsFs())
+		formatter = cliui.NewOutputFormatter(
+			cliui.TableFormat([]statsRow{}, []string{
+				"host_cpu",
+				"host_memory",
+				"home_disk",
+				"container_cpu",
+				"container_memory",
+			}),
+			cliui.JSONFormat(),
+		)
+	)
 	cmd := &clibase.Cmd{
-		Use:   "stat",
-		Short: "Show resource usage for the current workspace.",
+		Use:        "stat",
+		Short:      "Show resource usage for the current workspace.",
+		Middleware: initStatterMW(&st, fs),
 		Children: []*clibase.Cmd{
-			r.statCPU(st, fs),
-			r.statMem(st, fs),
-			r.statDisk(st),
+			r.statCPU(fs),
+			r.statMem(fs),
+			r.statDisk(fs),
 		},
 		Handler: func(inv *clibase.Invocation) error {
 			var sr statsRow
@@ -118,12 +130,16 @@ func (r *RootCmd) stat() *clibase.Cmd {
 	return cmd
 }
 
-func (*RootCmd) statCPU(s *clistat.Statter, fs afero.Fs) *clibase.Cmd {
-	var hostArg bool
-	formatter := cliui.NewOutputFormatter(cliui.TextFormat(), cliui.JSONFormat())
+func (*RootCmd) statCPU(fs afero.Fs) *clibase.Cmd {
+	var (
+		hostArg   bool
+		st        *clistat.Statter
+		formatter = cliui.NewOutputFormatter(cliui.TextFormat(), cliui.JSONFormat())
+	)
 	cmd := &clibase.Cmd{
-		Use:   "cpu",
-		Short: "Show CPU usage, in cores.",
+		Use:        "cpu",
+		Short:      "Show CPU usage, in cores.",
+		Middleware: initStatterMW(&st, fs),
 		Options: clibase.OptionSet{
 			{
 				Flag:        "host",
@@ -135,9 +151,9 @@ func (*RootCmd) statCPU(s *clistat.Statter, fs afero.Fs) *clibase.Cmd {
 			var cs *clistat.Result
 			var err error
 			if ok, _ := clistat.IsContainerized(fs); ok && !hostArg {
-				cs, err = s.ContainerCPU()
+				cs, err = st.ContainerCPU()
 			} else {
-				cs, err = s.HostCPU()
+				cs, err = st.HostCPU()
 			}
 			if err != nil {
 				return err
@@ -155,13 +171,17 @@ func (*RootCmd) statCPU(s *clistat.Statter, fs afero.Fs) *clibase.Cmd {
 	return cmd
 }
 
-func (*RootCmd) statMem(s *clistat.Statter, fs afero.Fs) *clibase.Cmd {
-	var hostArg bool
-	var prefixArg string
-	formatter := cliui.NewOutputFormatter(cliui.TextFormat(), cliui.JSONFormat())
+func (*RootCmd) statMem(fs afero.Fs) *clibase.Cmd {
+	var (
+		hostArg   bool
+		prefixArg string
+		st        *clistat.Statter
+		formatter = cliui.NewOutputFormatter(cliui.TextFormat(), cliui.JSONFormat())
+	)
 	cmd := &clibase.Cmd{
-		Use:   "mem",
-		Short: "Show memory usage, in gigabytes.",
+		Use:        "mem",
+		Short:      "Show memory usage, in gigabytes.",
+		Middleware: initStatterMW(&st, fs),
 		Options: clibase.OptionSet{
 			{
 				Flag:        "host",
@@ -185,9 +205,9 @@ func (*RootCmd) statMem(s *clistat.Statter, fs afero.Fs) *clibase.Cmd {
 			var ms *clistat.Result
 			var err error
 			if ok, _ := clistat.IsContainerized(fs); ok && !hostArg {
-				ms, err = s.ContainerMemory(pfx)
+				ms, err = st.ContainerMemory(pfx)
 			} else {
-				ms, err = s.HostMemory(pfx)
+				ms, err = st.HostMemory(pfx)
 			}
 			if err != nil {
 				return err
@@ -205,13 +225,17 @@ func (*RootCmd) statMem(s *clistat.Statter, fs afero.Fs) *clibase.Cmd {
 	return cmd
 }
 
-func (*RootCmd) statDisk(s *clistat.Statter) *clibase.Cmd {
-	var pathArg string
-	var prefixArg string
-	formatter := cliui.NewOutputFormatter(cliui.TextFormat(), cliui.JSONFormat())
+func (*RootCmd) statDisk(fs afero.Fs) *clibase.Cmd {
+	var (
+		pathArg   string
+		prefixArg string
+		st        *clistat.Statter
+		formatter = cliui.NewOutputFormatter(cliui.TextFormat(), cliui.JSONFormat())
+	)
 	cmd := &clibase.Cmd{
-		Use:   "disk",
-		Short: "Show disk usage, in gigabytes.",
+		Use:        "disk",
+		Short:      "Show disk usage, in gigabytes.",
+		Middleware: initStatterMW(&st, fs),
 		Options: clibase.OptionSet{
 			{
 				Flag:        "path",
@@ -237,7 +261,7 @@ func (*RootCmd) statDisk(s *clistat.Statter) *clibase.Cmd {
 			if len(inv.Args) > 0 {
 				pathArg = inv.Args[0]
 			}
-			ds, err := s.Disk(pfx, pathArg)
+			ds, err := st.Disk(pfx, pathArg)
 			if err != nil {
 				if os.IsNotExist(err) {
 					//nolint:gocritic // fmt.Errorf produces a more concise error.