coder · ammario · Mar 29, 2023 · Mar 29, 2023 · matifali · Apr 28, 2023
diff --git a/cli/root.go b/cli/root.go
@@ -106,6 +106,7 @@ func (r *RootCmd) Core() []*clibase.Cmd {
 		r.scaletest(),
 		r.gitssh(),
 		r.vscodeSSH(),
+		r.stat(),
 	}
 }
 
@@ -679,6 +680,15 @@ func (r *RootCmd) checkVersions(i *clibase.Invocation, client *codersdk.Client)
 	return nil
 }
 
+// verboseStderr returns the stderr writer if verbose is set, otherwise
+// it returns a discard writer.
+func (r *RootCmd) verboseStderr(inv *clibase.Invocation) io.Writer {
+	if r.verbose {
+		return inv.Stderr
+	}
+	return io.Discard
+}
+
 func (r *RootCmd) checkWarnings(i *clibase.Invocation, client *codersdk.Client) error {
 	if r.noFeatureWarning {
 		return nil

diff --git a/cli/stat.go b/cli/stat.go
@@ -0,0 +1,170 @@
+package cli
+
+import (
+	"bufio"
+	"bytes"
+	"fmt"
+	"os"
+	"strconv"
+	"strings"
+	"time"
+
+	"github.com/shirou/gopsutil/cpu"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/coder/cli/clibase"
+	"github.com/coder/coder/cli/cliui"
+)
+
+type statCmd struct {
+	*RootCmd
+	watch time.Duration
+}
+
+func (r *RootCmd) stat() *clibase.Cmd {
+	c := &clibase.Cmd{
+		Use:   "stat <type> [flags...]",
+		Short: "Display local system resource usage statistics",
+		Long:  "stat calls can be used as the script for agent metadata blocks.",
+	}
+	sc := statCmd{RootCmd: r}
+	c.Options.Add(
+		clibase.Option{
+			Flag:          "watch",
+			FlagShorthand: "w",
+			Description:   "Continuously display the statistic on the given interval.",
+			Value:         clibase.DurationOf(&sc.watch),
+		},
+	)
+	c.AddSubcommands(
+		sc.cpu(),
+	)
+	sc.setWatchLoops(c)
+	return c
+}
+
+func (sc *statCmd) setWatchLoops(c *clibase.Cmd) {
+	for _, cmd := range c.Children {
+		innerHandler := cmd.Handler
+		cmd.Handler = func(inv *clibase.Invocation) error {
+			if sc.watch == 0 {
+				return innerHandler(inv)
+			}
+
+			ticker := time.NewTicker(sc.watch)
+			defer ticker.Stop()
+
+			for range ticker.C {
+				if err := innerHandler(inv); err != nil {
+					_, _ = fmt.Fprintf(inv.Stderr, "error: %v", err)
+				}
+			}
+			panic("unreachable")
+		}
+	}
+}
+
+func cpuUsageFromCgroup(interval time.Duration) (float64, error) {
+	cgroup, err := os.OpenFile("/proc/self/cgroup", os.O_RDONLY, 0)
+	if err != nil {
+		return 0, err
+	}
+	defer cgroup.Close()
+	sc := bufio.NewScanner(cgroup)
+
+	var groupDir string
+	for sc.Scan() {
+		fields := strings.Split(sc.Text(), ":")
+		if len(fields) != 3 {
+			continue
+		}
+		if fields[1] != "cpu,cpuacct" {
+			continue
+		}
+		groupDir = fields[2]
+		break
+	}
+
+	if groupDir == "" {
+		return 0, xerrors.New("no cpu cgroup found")
+	}
+
+	cpuAcct := func() (int64, error) {
+		path := fmt.Sprintf("/sys/fs/cgroup/cpu,cpuacct/%s/cpuacct.usage", groupDir)
+
+		byt, err := os.ReadFile(
+			path,
+		)
+		if err != nil {
+			return 0, err
+		}
+
+		return strconv.ParseInt(string(bytes.TrimSpace(byt)), 10, 64)
+	}
+
+	stat1, err := cpuAcct()
+	if err != nil {
+		return 0, err
+	}
+
+	time.Sleep(interval)
+
+	stat2, err := cpuAcct()
+	if err != nil {
+		return 0, err
+	}
+
+	var (
+		cpuTime  = time.Duration(stat2 - stat1)
+		realTime = interval
+	)
+
+	ncpu, err := cpu.Counts(true)
+	if err != nil {
+		return 0, err
+	}
+
+	return (cpuTime.Seconds() / realTime.Seconds()) * 100 / float64(ncpu), nil
+}
+
+//nolint:revive
+func (sc *statCmd) cpu() *clibase.Cmd {
+	var interval time.Duration
+	c := &clibase.Cmd{
+		Use:     "cpu-usage",
+		Aliases: []string{"cu"},
+		Short:   "Display the system's cpu usage",
+		Long:    "If inside a cgroup (e.g. docker container), the cpu usage is ",
+		Handler: func(inv *clibase.Invocation) error {
+			if sc.watch != 0 {
+				interval = sc.watch
+			}
+
+			r, err := cpuUsageFromCgroup(interval)
+			if err != nil {
+				cliui.Infof(sc.verboseStderr(inv), "cgroup error: %+v", err)
+
+				// Use standard methods if cgroup method fails.
+				rs, err := cpu.Percent(interval, false)
+				if err != nil {
+					return err
+				}
+				r = rs[0]
+			}
+
+			_, _ = fmt.Fprintf(inv.Stdout, "%02.0f\n", r)
+
+			return nil
+		},
+		Options: []clibase.Option{
+			{
+				Flag:          "interval",
+				FlagShorthand: "i",
+				Description:   `The sample collection interval. If --watch is set, it overrides this value.`,
+				Default:       "0s",
+				Value:         clibase.DurationOf(&interval),
+			},
+		},
+	}
+	return c
+}
diff --git a/cli/stat_test.go b/cli/stat_test.go
@@ -0,0 +1,24 @@
+package cli_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/coder/coder/cli/clitest"
+)
+
+func TestStat(t *testing.T) {
+	t.Parallel()
+
+	t.Run("cpu", func(t *testing.T) {
+		t.Parallel()
+		inv, _ := clitest.New(t, "stat", "cpu")
+		var out bytes.Buffer
+		inv.Stdout = &out
+		clitest.Run(t, inv)
+
+		require.Regexp(t, `^[\d]{2}\n$`, out.String())
+	})
+}
diff --git a/cli/testdata/coder_--help.golden b/cli/testdata/coder_--help.golden
@@ -34,6 +34,7 @@ Coder v0.0.0-devel — A tool for provisioning self-hosted development environme
                       workspace
     ssh               Start a shell into a workspace
     start             Start a workspace
+    stat              Display local system resource usage statistics
     state             Manually manage Terraform state to fix broken workspaces
     stop              Stop a workspace
     templates         Manage templates

diff --git a/cli/testdata/coder_stat_--help.golden b/cli/testdata/coder_stat_--help.golden
@@ -0,0 +1,15 @@
+Usage: coder stat [flags] <type> [flags...]
+
+Display local system resource usage statistics
+
+stat calls can be used as the script for agent metadata blocks.
+
+[1mSubcommands[0m
+    cpu    Display the system's cpu usage
+
+[1mOptions[0m
+  -w, --watch duration
+          Continuously display the statistic on the given interval.
+
+---
+Run `coder --help` for a list of global options.
diff --git a/cli/testdata/coder_stat_cpu_--help.golden b/cli/testdata/coder_stat_cpu_--help.golden
@@ -0,0 +1,13 @@
+Usage: coder stat cpu [flags]
+
+Display the system's cpu usage
+
+If inside a cgroup (e.g. docker container), the cpu usage is
+
+[1mOptions[0m
+  -i, --interval duration (default: 0s)
+          The sample collection interval. If --watch is set, it overrides this
+          value.
+
+---
+Run `coder --help` for a list of global options.
diff --git a/docs/cli.md b/docs/cli.md
@@ -49,6 +49,7 @@ Coder — A tool for provisioning self-hosted development environments with Terr
 | [<code>speedtest</code>](./cli/speedtest)           | Run upload and download tests from your machine to a workspace         |
 | [<code>ssh</code>](./cli/ssh)                       | Start a shell into a workspace                                         |
 | [<code>start</code>](./cli/start)                   | Start a workspace                                                      |
+| [<code>stat</code>](./cli/stat)                     | Display local system resource usage statistics                         |
 | [<code>state</code>](./cli/state)                   | Manually manage Terraform state to fix broken workspaces               |
 | [<code>stop</code>](./cli/stop)                     | Stop a workspace                                                       |
 | [<code>templates</code>](./cli/templates)           | Manage templates                                                       |

diff --git a/docs/cli/stat.md b/docs/cli/stat.md
@@ -0,0 +1,33 @@
+<!-- DO NOT EDIT | GENERATED CONTENT -->
+
+# stat
+
+Display local system resource usage statistics
+
+## Usage
+
+```console
+coder stat [flags] <type> [flags...]
+```
+
+## Description
+
+```console
+stat calls can be used as the script for agent metadata blocks.
+```
+
+## Subcommands
+
+| Name                           | Purpose                        |
+| ------------------------------ | ------------------------------ |
+| [<code>cpu</code>](./stat_cpu) | Display the system's cpu usage |
+
+## Options
+
+### -w, --watch
+
+|      |                       |
+| ---- | --------------------- |
+| Type | <code>duration</code> |
+
+Continuously display the statistic on the given interval.
diff --git a/docs/cli/stat_cpu.md b/docs/cli/stat_cpu.md
@@ -0,0 +1,28 @@
+<!-- DO NOT EDIT | GENERATED CONTENT -->
+
+# stat cpu
+
+Display the system's cpu usage
+
+## Usage
+
+```console
+coder stat cpu [flags]
+```
+
+## Description
+
+```console
+If inside a cgroup (e.g. docker container), the cpu usage is
+```
+
+## Options
+
+### -i, --interval
+
+|         |                       |
+| ------- | --------------------- |
+| Type    | <code>duration</code> |
+| Default | <code>0s</code>       |
+
+The sample collection interval. If --watch is set, it overrides this value.
diff --git a/docs/manifest.json b/docs/manifest.json
@@ -677,6 +677,16 @@
           "description": "Start a workspace",
           "path": "cli/start.md"
         },
+        {
+          "title": "stat",
+          "description": "Display local system resource usage statistics",
+          "path": "cli/stat.md"
+        },
+        {
+          "title": "stat cpu",
+          "description": "Display the system's cpu usage",
+          "path": "cli/stat_cpu.md"
+        },
         {
           "title": "state",
           "description": "Manually manage Terraform state to fix broken workspaces",

diff --git a/go.mod b/go.mod
@@ -127,6 +127,7 @@ require (
 	github.com/prometheus/client_golang v1.14.0
 	github.com/quasilyte/go-ruleguard/dsl v0.3.21
 	github.com/robfig/cron/v3 v3.0.1
+	github.com/shirou/gopsutil v3.21.11+incompatible
 	github.com/spf13/afero v1.9.3
 	github.com/spf13/pflag v1.0.5
 	github.com/stretchr/testify v1.8.1
@@ -187,6 +188,9 @@ require (
 	github.com/modern-go/reflect2 v1.0.2 // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
 	github.com/swaggo/files/v2 v2.0.0 // indirect
+	github.com/tklauser/go-sysconf v0.3.9 // indirect
+	github.com/tklauser/numcpus v0.3.0 // indirect
+	github.com/yusufpapurcu/wmi v1.2.2 // indirect
 	golang.org/x/text v0.8.0 // indirect
 	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230215201556-9c5414ab4bde // indirect
 )