chultun/internal/budget/proc_linux.go

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
}