Add systemd resource budgets with TOML profiles

internal/budget/proc_test.go

@@ -0,0 +1,111 @@
+package budget
+
+import (
+	"os"
+	"path/filepath"
+	"reflect"
+	"testing"
+)
+
+func TestClassify(t *testing.T) {
+	c := testConfig()
+	terminal := "/user.slice/app.slice/app-terminal-surface-123.scope"
+	job := "/user.slice/app.slice/budget-c123abc.scope"
+	for _, tc := range []struct{ exe, group, profile, slice string }{
+		{"rg", terminal, "search", ""}, {"java", terminal, "compute", ""},
+		{"python3.14", terminal, "compute", ""}, {"python", terminal, "compute", ""},
+		{"rsync", terminal, "bulk", ""}, {"editor", terminal, "default", "app-editor.slice"},
+		{"cargo", terminal, "", ""}, {"kubectl", terminal, "", ""}, {"codex", terminal, "", ""},
+		{"npm", terminal, "", ""}, {"make", terminal, "", ""}, {"dotnet", terminal, "", ""},
+		{"rg", job, "search", ""}, {"java", job, "", ""}, {"python3", job, "", ""},
+		{"editor", job, "default", "app-editor.slice"},
+		{"rg", "/user.slice/app.slice/work-123.scope", "", ""},
+		{"rg", "/user.slice/app.slice/app-flatpak-browser-123.scope", "", ""},
+		{"rg", "/system.slice/docker-123.scope", "", ""},
+		{"java", "/user.slice/app.slice/containerd.service/container-123", "", ""},
+		{"rg", "/user.slice/search.slice/budget-s123abc.scope", "", ""},
+	} {
+		t.Run(tc.exe+tc.group, func(t *testing.T) {
+			profile, slice := c.classify(tc.exe, tc.group)
+			if profile != tc.profile || slice != tc.slice {
+				t.Fatalf("got %q, %q; want %q, %q", profile, slice, tc.profile, tc.slice)
+			}
+		})
+	}
+}
+
+func TestProcParent(t *testing.T) {
+	for _, stat := range []string{"123 (rg) R 42 0 0", "123 (a ) b (c)) S 42 0 0", "123 (a\nb) S 42 0 0"} {
+		got, err := parseParent(stat)
+		if err != nil || got != 42 {
+			t.Fatalf("%q: %d, %v", stat, got, err)
+		}
+	}
+	for _, stat := range []string{"", "123 (rg) ", "123 (rg) S bad"} {
+		if _, err := parseParent(stat); err == nil {
+			t.Errorf("accepted %q", stat)
+		}
+	}
+}
+
+func put(t *testing.T, path, text string) {
+	t.Helper()
+	if err := os.MkdirAll(filepath.Dir(path), 0755); err != nil {
+		t.Fatal(err)
+	}
+	if err := os.WriteFile(path, []byte(text), 0644); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestProcSnapshot(t *testing.T) {
+	root := t.TempDir()
+	group := "/app.slice/app-terminal-surface-123.scope"
+	for _, pid := range []string{"100", "101"} {
+		put(t, filepath.Join(root, pid, "cgroup"), "0::"+group+"\n")
+		put(t, filepath.Join(root, pid, "stat"), pid+" (rg ) test) S 10 0 0\n")
+	}
+	if err := os.Symlink("/opt/bundled/rg", filepath.Join(root, "100/exe")); err != nil {
+		t.Fatal(err)
+	}
+	put(t, filepath.Join(root, "200/cgroup"), "0::/system.slice/docker-123.scope\n")
+	if err := os.Mkdir(filepath.Join(root, "self"), 0755); err != nil {
+		t.Fatal(err)
+	}
+	p := procFS{root, uint32(os.Getuid())}
+	tasks, err := p.Snapshot(testConfig())
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := map[uint32]task{100: {10, group, "search", ""}}
+	if !reflect.DeepEqual(tasks, want) {
+		t.Fatalf("got %#v, want %#v", tasks, want)
+	}
+	got, err := p.Group(100)
+	if err != nil || got != group {
+		t.Fatalf("%q, %v", got, err)
+	}
+	p.uid++
+	tasks, err = p.Snapshot(testConfig())
+	if err != nil || len(tasks) != 0 {
+		t.Fatalf("other UID included: %v, %v", tasks, err)
+	}
+	p.root = filepath.Join(root, "missing")
+	if _, err := p.Snapshot(testConfig()); !os.IsNotExist(err) {
+		t.Fatal(err)
+	}
+}
+
+func TestFamily(t *testing.T) {
+	tasks := map[uint32]task{
+		10: {Parent: 1, Group: "/terminal"},
+		11: {Parent: 10, Group: "/terminal"},
+		12: {Parent: 11, Group: "/terminal"},
+		13: {Parent: 10, Group: "/container"},
+		14: {Parent: 13, Group: "/terminal"},
+		15: {Parent: 1, Group: "/terminal"},
+	}
+	if got := family(tasks, 10); !reflect.DeepEqual(got, []uint32{10, 11, 12}) {
+		t.Fatal(got)
+	}
+}