Add systemd resource budgets with TOML profiles

internal/budget/config.go

@@ -0,0 +1,168 @@
+package budget
+
+import (
+	"fmt"
+	"maps"
+	"math"
+	"slices"
+	"strconv"
+	"strings"
+
+	"github.com/BurntSushi/toml"
+	"github.com/godbus/dbus/v5"
+)
+
+type profile struct {
+	CPU    string `toml:"cpu"`
+	High   string `toml:"high"`
+	Memory string `toml:"memory"`
+	Swap   string `toml:"swap"`
+}
+
+type ioLimit struct {
+	Read  string `toml:"read"`
+	Write string `toml:"write"`
+	RIOPS string `toml:"riops"`
+	WIOPS string `toml:"wiops"`
+}
+
+type device struct {
+	Path   string             `toml:"path"`
+	Limits map[string]ioLimit `toml:"limits"`
+}
+
+type config struct {
+	Profiles  map[string]profile  `toml:"profiles"`
+	Devices   []device            `toml:"devices"`
+	Apps      map[string][]string `toml:"apps"`
+	Units     map[string][]string `toml:"units"`
+	Prefixes  map[string][]string `toml:"prefixes"`
+	Commands  map[string]string   `toml:"commands"`
+	Python    string              `toml:"python"`
+	Terminals []string            `toml:"terminals"`
+	Shared    map[string]string   `toml:"shared"`
+}
+
+type property struct {
+	Name  string
+	Value dbus.Variant
+}
+
+type diskLimit struct {
+	Path  string
+	Limit uint64
+}
+
+type ioField struct {
+	name string
+	get  func(ioLimit) string
+}
+
+var ioFields = []ioField{
+	{"IOReadBandwidthMax", func(l ioLimit) string { return l.Read }},
+	{"IOWriteBandwidthMax", func(l ioLimit) string { return l.Write }},
+	{"IOReadIOPSMax", func(l ioLimit) string { return l.RIOPS }},
+	{"IOWriteIOPSMax", func(l ioLimit) string { return l.WIOPS }},
+}
+
+func keys[V any](m map[string]V) []string {
+	return slices.Sorted(maps.Keys(m))
+}
+
+func loadConfig(path string) (config, error) {
+	var c config
+	_, err := toml.DecodeFile(path, &c)
+	return c, err
+}
+
+func parseSize(value string, base uint64) (uint64, error) {
+	if value == "infinity" {
+		return math.MaxUint64, nil
+	}
+	factor := uint64(1)
+	for i, suffix := range []string{"K", "M", "G"} {
+		if strings.HasSuffix(value, suffix) {
+			value = strings.TrimSuffix(value, suffix)
+			for range i + 1 {
+				factor *= base
+			}
+			break
+		}
+	}
+	number, err := strconv.ParseUint(value, 10, 64)
+	return number * factor, err
+}
+
+func prop(name string, value any) property {
+	return property{Name: name, Value: dbus.MakeVariant(value)}
+}
+
+func (c config) getProfile(name string) (profile, error) {
+	p, ok := c.Profiles[name]
+	if !ok {
+		return profile{}, fmt.Errorf("unknown profile %q", name)
+	}
+	return p, nil
+}
+
+func (c config) properties(name string) ([]property, error) {
+	p, err := c.getProfile(name)
+	if err != nil {
+		return nil, err
+	}
+	quota := uint64(math.MaxUint64)
+	if p.CPU != "" {
+		quota, err = strconv.ParseUint(strings.TrimSuffix(p.CPU, "%"), 10, 64)
+		if err != nil {
+			return nil, fmt.Errorf("CPUQuota: %w", err)
+		}
+		quota *= 10000
+	}
+	result := []property{prop("CPUQuotaPerSecUSec", quota)}
+	for _, field := range []struct{ name, value string }{
+		{"MemoryHigh", p.High}, {"MemoryMax", p.Memory}, {"MemorySwapMax", p.Swap},
+	} {
+		value, err := parseSize(field.value, 1024)
+		if err != nil {
+			return nil, fmt.Errorf("%s: %w", field.name, err)
+		}
+		result = append(result, prop(field.name, value))
+	}
+	for _, field := range ioFields {
+		limits := make([]diskLimit, 0, len(c.Devices))
+		for _, d := range c.Devices {
+			limit, ok := d.Limits[name]
+			if !ok {
+				continue
+			}
+			value, err := parseSize(field.get(limit), 1000)
+			if err != nil {
+				return nil, fmt.Errorf("%s: %w", field.name, err)
+			}
+			limits = append(limits, diskLimit{Path: d.Path, Limit: value})
+		}
+		result = append(result, prop(field.name, limits))
+	}
+	return result, nil
+}
+
+func (c config) assignments(name string) ([]string, error) {
+	p, err := c.getProfile(name)
+	if err != nil {
+		return nil, err
+	}
+	result := []string{"CPUQuota=" + p.CPU, "MemoryHigh=" + p.High, "MemoryMax=" + p.Memory, "MemorySwapMax=" + p.Swap}
+	for _, field := range ioFields {
+		found := false
+		for _, d := range c.Devices {
+			if limit, ok := d.Limits[name]; ok {
+				result = append(result, field.name+"="+d.Path+" "+field.get(limit))
+				found = true
+			}
+		}
+		if !found {
+			result = append(result, field.name+"=")
+		}
+	}
+	return result, nil
+}