chultun/internal/budget/config.go

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)
	if err != nil {
		return c, err
	}
	return c, c.validate()
}

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 != "" {
		percentage, ok := strings.CutSuffix(p.CPU, "%")
		if !ok {
			return nil, fmt.Errorf("CPUQuota: expected a percentage, got %q", p.CPU)
		}
		quota, err = strconv.ParseUint(percentage, 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
}