Add systemd resource budgets with TOML profiles

2026-10-05 19:19:05 +09:00 · 2026-10-05 19:19:05 +09:00 · be8749a384
commit be8749a384
1 changed files with 2591 additions and 0 deletions
--- a/internal/budget/watch.go
+++ b/internal/budget/watch.go
@ -0,0 +1,153 @@
+package budget
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"slices"
+	"strings"
+	"time"
+)
+
+type watcher struct {
+	config config
+	ctl    unitCtl
+	procs  procSource
+	cgroup string
+	owned  map[string]bool
+	out    io.Writer
+}
+
+func randomID() (string, error) {
+	var data [3]byte
+	_, err := rand.Read(data[:])
+	return hex.EncodeToString(data[:]), err
+}
+
+func scopeProps(c config, name, slice string, pids []uint32) ([]property, error) {
+	if slice == "" {
+		slice = "app.slice"
+		if name == "search" {
+			slice = "search.slice"
+		}
+	}
+	values := []property{prop("Description", name+" workload"), prop("Slice", slice), prop("CollectMode", "inactive-or-failed"), prop("PIDs", pids)}
+	if name != "search" {
+		limits, err := c.properties(name)
+		if err != nil {
+			return nil, err
+		}
+		values = append(values, limits...)
+	}
+	return values, nil
+}
+
+func (w *watcher) init() error {
+	units, err := w.ctl.Units()
+	if err != nil {
+		return err
+	}
+	w.owned = map[string]bool{}
+	for _, unit := range units {
+		if strings.HasPrefix(unit.Name, "budget-") && strings.HasSuffix(unit.Name, ".scope") {
+			w.owned[unit.Name] = true
+		}
+	}
+	return nil
+}
+
+func (w *watcher) reap() error {
+	for _, name := range keys(w.owned) {
+		state, err := w.ctl.State(name)
+		if busError(err, noUnit) {
+			delete(w.owned, name)
+			continue
+		}
+		if err != nil {
+			return err
+		}
+		switch state.Active {
+		case "inactive", "failed":
+			delete(w.owned, name)
+		case "active":
+			data, err := os.ReadFile(filepath.Join(w.cgroup, strings.TrimPrefix(state.Group, "/"), "cgroup.events"))
+			if os.IsNotExist(err) {
+				delete(w.owned, name)
+				continue
+			}
+			if err != nil {
+				return err
+			}
+			if strings.Contains(string(data), "populated 0") {
+				if err := w.ctl.Stop(name); err != nil && !busError(err, noUnit) {
+					return err
+				}
+				delete(w.owned, name)
+			}
+		}
+	}
+	return nil
+}
+
+func (w *watcher) step() error {
+	tasks, err := w.procs.Snapshot(w.config)
+	if err != nil {
+		return err
+	}
+	pids := make([]uint32, 0, len(tasks))
+	for pid := range tasks {
+		pids = append(pids, pid)
+	}
+	slices.Sort(pids)
+	for _, pid := range pids {
+		value := tasks[pid]
+		if value.Profile == "" {
+			continue
+		}
+		group, err := w.procs.Group(pid)
+		if gone(err) {
+			continue
+		}
+		if err != nil {
+			return err
+		}
+		if group != value.Group {
+			continue
+		}
+		members := family(tasks, pid)
+		values, err := scopeProps(w.config, value.Profile, value.Slice, members)
+		if err != nil {
+			return err
+		}
+		id, err := randomID()
+		if err != nil {
+			return err
+		}
+		name := "budget-" + value.Profile[:1] + id + ".scope"
+		if err = w.ctl.Start(name, values); busError(err, noPID) {
+			continue
+		} else if err != nil {
+			return err
+		}
+		w.owned[name] = true
+		if _, err := fmt.Fprintf(w.out, "%s: %s %v\n", value.Profile, name, members); err != nil {
+			return err
+		}
+	}
+	return w.reap()
+}
+
+func (w *watcher) watch() error {
+	if err := w.init(); err != nil {
+		return err
+	}
+	for {
+		if err := w.step(); err != nil {
+			return err
+		}
+		time.Sleep(200 * time.Millisecond)
+	}
+}