Add systemd resource budgets with TOML profiles

internal/budget/cli.go

@@ -0,0 +1,103 @@
+package budget
+
+import (
+	"errors"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/BurntSushi/toml"
+)
+
+const usage = `usage: budget COMMAND
+
+  apply                    Apply configured limits and restart the watcher
+  watch                    Assign profiles to terminal workloads
+  status                   Show the configured resource profiles
+  run PROFILE -- COMMAND   Launch a native command in a profile
+  set PROFILE TARGET       Override a running app ID or systemd unit
+
+Configuration: ~/.config/budget/profiles.toml
+`
+
+type request struct {
+	action  string
+	profile string
+	target  string
+	command []string
+}
+
+func parseArgs(args []string) (request, error) {
+	if len(args) == 0 {
+		return request{}, fmt.Errorf("a command is required")
+	}
+	r := request{action: args[0]}
+	switch r.action {
+	case "-h", "--help":
+		r.action = "help"
+	case "apply", "watch", "status":
+		if len(args) != 1 {
+			return request{}, fmt.Errorf("%s takes no arguments", r.action)
+		}
+	case "run":
+		if len(args) < 3 {
+			return request{}, fmt.Errorf("usage: budget run PROFILE -- COMMAND")
+		}
+		r.profile, r.command = args[1], args[2:]
+	case "set":
+		if len(args) != 3 {
+			return request{}, fmt.Errorf("usage: budget set PROFILE TARGET")
+		}
+		r.profile, r.target = args[1], args[2]
+	default:
+		return request{}, fmt.Errorf("unknown command %q", r.action)
+	}
+	return r, nil
+}
+
+func Run(args []string, out io.Writer) (err error) {
+	r, err := parseArgs(args)
+	if err != nil {
+		return err
+	}
+	if r.action == "help" {
+		_, err := io.WriteString(out, usage)
+		return err
+	}
+	home, err := os.UserHomeDir()
+	if err != nil {
+		return err
+	}
+	paths := homePaths(home)
+	c, err := loadConfig(paths.config)
+	if err != nil {
+		return err
+	}
+	if r.profile != "" {
+		if _, err := c.getProfile(r.profile); err != nil {
+			return err
+		}
+	}
+	if r.action == "status" {
+		return toml.NewEncoder(out).Encode(c)
+	}
+	if r.action == "run" {
+		return launch(c, r.profile, r.command)
+	}
+	m, err := connect()
+	if err != nil {
+		return err
+	}
+	defer func() { err = errors.Join(err, m.conn.Close()) }()
+	a := application{config: c, paths: paths, ctl: m, run: systemctl, out: out}
+	switch r.action {
+	case "apply":
+		return a.apply()
+	case "set":
+		return a.override(r.profile, r.target)
+	case "watch":
+		w := watcher{config: c, ctl: m, procs: procFS{root: "/proc", uid: uint32(os.Getuid())}, cgroup: "/sys/fs/cgroup", out: out}
+		return w.watch()
+	}
+	return nil
+}