Add systemd resource budgets with TOML profiles

internal/budget/proc_linux.go

@@ -0,0 +1,169 @@
+package budget
+
+import (
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
+	"regexp"
+	"slices"
+	"strconv"
+	"strings"
+	"syscall"
+)
+
+type task struct {
+	Parent  uint32
+	Group   string
+	Profile string
+	Slice   string
+}
+
+type procSource interface {
+	Snapshot(config) (map[uint32]task, error)
+	Group(uint32) (string, error)
+}
+
+type procFS struct {
+	root string
+	uid  uint32
+}
+
+var pythonName = regexp.MustCompile(`^python[23]?(?:\.\d+)?$`)
+
+func gone(err error) bool {
+	return errors.Is(err, os.ErrNotExist) || errors.Is(err, syscall.ESRCH)
+}
+
+func (c config) sourceGroup(group string) (bool, bool) {
+	unit := filepath.Base(group)
+	terminal := false
+	for _, prefix := range c.Terminals {
+		if strings.HasPrefix(unit, prefix) {
+			terminal = true
+			break
+		}
+	}
+	job := (strings.HasPrefix(unit, "budget-c") || strings.HasPrefix(unit, "budget-b")) && strings.HasSuffix(unit, ".scope")
+	return terminal, job
+}
+
+func (c config) classify(exe, group string) (string, string) {
+	terminal, job := c.sourceGroup(group)
+	if !terminal && !job {
+		return "", ""
+	}
+	name := c.Commands[exe]
+	slice := ""
+	if shared, ok := c.Shared[exe]; ok {
+		name, slice = shared, c.sharedSlice()
+	}
+	if pythonName.MatchString(exe) {
+		name = c.Python
+	}
+	if job && name != "search" && slice == "" {
+		name = ""
+	}
+	return name, slice
+}
+
+func readGroup(path string) (string, error) {
+	data, err := os.ReadFile(path)
+	return strings.TrimPrefix(strings.TrimSpace(string(data)), "0::"), err
+}
+
+func (p procFS) Group(pid uint32) (string, error) {
+	return readGroup(filepath.Join(p.root, strconv.FormatUint(uint64(pid), 10), "cgroup"))
+}
+
+func parseParent(stat string) (uint32, error) {
+	i := strings.LastIndex(stat, ") ")
+	if i < 0 {
+		return 0, fmt.Errorf("invalid proc stat")
+	}
+	fields := strings.Fields(stat[i+2:])
+	if len(fields) < 2 {
+		return 0, fmt.Errorf("missing parent in proc stat")
+	}
+	value, err := strconv.ParseUint(fields[1], 10, 32)
+	return uint32(value), err
+}
+
+func (p procFS) readTask(path string, c config) (task, bool, error) {
+	info, err := os.Stat(path)
+	if err != nil {
+		return task{}, false, err
+	}
+	if info.Sys().(*syscall.Stat_t).Uid != p.uid {
+		return task{}, false, nil
+	}
+	group, err := readGroup(filepath.Join(path, "cgroup"))
+	if err != nil {
+		return task{}, false, err
+	}
+	terminal, job := c.sourceGroup(group)
+	if !terminal && !job {
+		return task{}, false, nil
+	}
+	exe, err := os.Readlink(filepath.Join(path, "exe"))
+	if err != nil {
+		return task{}, false, err
+	}
+	stat, err := os.ReadFile(filepath.Join(path, "stat"))
+	if err != nil {
+		return task{}, false, err
+	}
+	parent, err := parseParent(string(stat))
+	if err != nil {
+		return task{}, false, err
+	}
+	name, slice := c.classify(filepath.Base(exe), group)
+	return task{Parent: parent, Group: group, Profile: name, Slice: slice}, true, nil
+}
+
+func (p procFS) Snapshot(c config) (map[uint32]task, error) {
+	entries, err := os.ReadDir(p.root)
+	if err != nil {
+		return nil, err
+	}
+	tasks := map[uint32]task{}
+	for _, entry := range entries {
+		pid, err := strconv.ParseUint(entry.Name(), 10, 32)
+		if err != nil {
+			continue
+		}
+		value, include, err := p.readTask(filepath.Join(p.root, entry.Name()), c)
+		if gone(err) {
+			continue
+		}
+		if err != nil {
+			return nil, err
+		}
+		if include {
+			tasks[uint32(pid)] = value
+		}
+	}
+	return tasks, nil
+}
+
+func family(tasks map[uint32]task, pid uint32) []uint32 {
+	group := tasks[pid].Group
+	seen := map[uint32]bool{pid: true}
+	pending := []uint32{pid}
+	for len(pending) > 0 {
+		parent := pending[0]
+		pending = pending[1:]
+		for child, value := range tasks {
+			if value.Parent == parent && value.Group == group && !seen[child] {
+				seen[child] = true
+				pending = append(pending, child)
+			}
+		}
+	}
+	result := make([]uint32, 0, len(seen))
+	for child := range seen {
+		result = append(result, child)
+	}
+	slices.Sort(result)
+	return result
+}