103 lines
2.4 KiB
Go
103 lines
2.4 KiB
Go
|
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
|
||
|
|
}
|