Add systemd resource budgets with TOML profiles

internal/budget/watch_test.go

@@ -0,0 +1,174 @@
+package budget
+
+import (
+	"errors"
+	"io"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/godbus/dbus/v5"
+)
+
+func TestSearchScope(t *testing.T) {
+	props, err := scopeProps(testConfig(), "search", "", []uint32{10, 11})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := values(props)
+	if got["Slice"].Value() != "search.slice" {
+		t.Fatal(got["Slice"])
+	}
+	if _, exists := got["MemoryMax"]; exists {
+		t.Fatal("search scopes should share the slice ceiling")
+	}
+	if got["PIDs"].Signature().String() != "au" || !reflect.DeepEqual(got["PIDs"].Value(), []uint32{10, 11}) {
+		t.Fatal(got["PIDs"])
+	}
+	props, err = scopeProps(testConfig(), "default", "app-editor.slice", []uint32{10})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got = values(props)
+	if got["Slice"].Value() != "app-editor.slice" || got["MemoryMax"].Value() != uint64(8589934592) {
+		t.Fatal(got)
+	}
+}
+
+func TestWatcherFamily(t *testing.T) {
+	group := "/app.slice/app-terminal-123.scope"
+	p := fakeProcs{map[uint32]task{10: {1, group, "compute", ""}, 11: {10, group, "search", ""}}, map[uint32]string{10: group, 11: group}}
+	f := newFake()
+	f.onStart = func(name string, props []property) {
+		for _, pid := range values(props)["PIDs"].Value().([]uint32) {
+			p.groups[pid] = "/app.slice/" + name
+		}
+	}
+	w := watcher{config: testConfig(), ctl: f, procs: p, cgroup: t.TempDir(), out: io.Discard}
+	if err := w.init(); err != nil {
+		t.Fatal(err)
+	}
+	if err := w.step(); err != nil {
+		t.Fatal(err)
+	}
+	if len(f.starts) != 1 {
+		t.Fatalf("family was assigned twice: %v", f.starts)
+	}
+	for name, props := range f.starts {
+		if len(name) > 20 || !strings.HasPrefix(name, "budget-c") {
+			t.Fatal(name)
+		}
+		if !reflect.DeepEqual(values(props)["PIDs"].Value(), []uint32{10, 11}) {
+			t.Fatal(props)
+		}
+	}
+	p.tasks = map[uint32]task{10: {1, p.groups[10], "", ""}, 11: {10, p.groups[11], "search", ""}}
+	w.procs = p
+	if err := w.step(); err != nil {
+		t.Fatal(err)
+	}
+	if len(f.starts) != 2 {
+		t.Fatal(f.starts)
+	}
+	for name, props := range f.starts {
+		if strings.HasPrefix(name, "budget-s") && values(props)["Slice"].Value() != "search.slice" {
+			t.Fatal(props)
+		}
+	}
+}
+
+func TestWatcherRaces(t *testing.T) {
+	group := "/app.slice/app-terminal-123.scope"
+	for _, tc := range []struct {
+		name, current string
+		err           error
+	}{
+		{"moved", "/other", nil}, {"exited", "", nil},
+		{"API exit", group, dbus.Error{Name: noPID}},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			p := fakeProcs{map[uint32]task{10: {1, group, "search", ""}}, map[uint32]string{}}
+			if tc.current != "" {
+				p.groups[10] = tc.current
+			}
+			f := newFake()
+			f.startErr = tc.err
+			w := watcher{config: testConfig(), ctl: f, procs: p, cgroup: t.TempDir(), out: io.Discard}
+			if err := w.init(); err != nil {
+				t.Fatal(err)
+			}
+			if err := w.step(); err != nil {
+				t.Fatal(err)
+			}
+			if len(w.owned) != 0 || len(f.starts) != 0 {
+				t.Fatal("an exited or moved task was retained")
+			}
+		})
+	}
+	f := newFake()
+	f.startErr = errors.New("bus unavailable")
+	w := watcher{config: testConfig(), ctl: f, procs: fakeProcs{map[uint32]task{10: {1, group, "search", ""}}, map[uint32]string{10: group}}, cgroup: t.TempDir(), owned: map[string]bool{}, out: io.Discard}
+	if err := w.step(); err != f.startErr {
+		t.Fatalf("unexpected error was hidden: %v", err)
+	}
+}
+
+func TestReap(t *testing.T) {
+	f := newFake()
+	f.states = map[string]unitState{
+		"budget-c000001.scope": {"active", "/empty"},
+		"budget-c000002.scope": {"active", "/live"},
+		"budget-c000003.scope": {"activating", ""},
+		"budget-c000004.scope": {"inactive", ""},
+		"budget-c000005.scope": {"failed", ""},
+		"budget-c000007.scope": {"active", "/removed"},
+	}
+	root := t.TempDir()
+	put(t, filepath.Join(root, "empty/cgroup.events"), "populated 0\nfrozen 0\n")
+	put(t, filepath.Join(root, "live/cgroup.events"), "populated 1\nfrozen 0\n")
+	w := watcher{ctl: f, cgroup: root, owned: map[string]bool{}}
+	for i := 1; i <= 7; i++ {
+		w.owned["budget-c00000"+string(rune('0'+i))+".scope"] = true
+	}
+	if err := w.reap(); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(f.stops, []string{"budget-c000001.scope"}) {
+		t.Fatal(f.stops)
+	}
+	if !reflect.DeepEqual(keys(w.owned), []string{"budget-c000002.scope", "budget-c000003.scope"}) {
+		t.Fatal(w.owned)
+	}
+	f.stateErr["budget-c000002.scope"] = errors.New("bus unavailable")
+	if err := w.reap(); err == nil {
+		t.Fatal("unexpected bus error hidden")
+	}
+}
+
+func TestRecoverScopes(t *testing.T) {
+	f := newFake()
+	f.units = []unit{{Name: "budget-c123abc.scope"}, {Name: "budget.service"}, {Name: "work-123.scope"}, {Name: "app-terminal-123.scope"}}
+	w := watcher{ctl: f}
+	if err := w.init(); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(keys(w.owned), []string{"budget-c123abc.scope"}) {
+		t.Fatal(w.owned)
+	}
+}
+
+func TestWatcherOutput(t *testing.T) {
+	want := errors.New("output failed")
+	group := "/app.slice/app-terminal-123.scope"
+	w := watcher{
+		config: testConfig(), ctl: newFake(), cgroup: t.TempDir(), out: failWriter{want},
+		procs: fakeProcs{map[uint32]task{10: {1, group, "search", ""}}, map[uint32]string{10: group}},
+	}
+	if err := w.init(); err != nil {
+		t.Fatal(err)
+	}
+	if err := w.step(); !errors.Is(err, want) {
+		t.Fatal(err)
+	}
+}