Fix profile updates and validate configuration before apply

internal/budget/app.go

@@ -30,7 +30,7 @@ func (a application) apply() error {
 		return err
 	}
 	for _, unit := range units {
-		name := a.config.unitProfile(unit.Name)
+		name := a.config.activeProfile(unit)
 		if name == "" || unit.Active != "active" {
 			continue
 		}

internal/budget/app_test.go

@@ -2,7 +2,10 @@ package budget
 
 import (
 	"errors"
+	"fmt"
 	"io"
+	"os"
+	"path/filepath"
 	"reflect"
 	"strings"
 	"testing"
@@ -46,6 +49,98 @@ func TestApply(t *testing.T) {
 	}
 }
 
+func TestApplyShared(t *testing.T) {
+	c := testConfig()
+	c.Shared["editor"] = "compute"
+	f := newFake()
+	f.units = []unit{{Name: "[email protected]", Active: "active"}}
+	a := application{c, homePaths(t.TempDir()), f, func(...string) error { return nil }, io.Discard}
+	if err := a.apply(); err != nil {
+		t.Fatal(err)
+	}
+	if got := values(f.sets[f.units[0].Name])["MemoryMax"].Value(); got != uint64(17179869184) {
+		t.Fatalf("shared app ceiling: got %v, want 16 GiB", got)
+	}
+	data, err := os.ReadFile(filepath.Join(a.paths.units, "[email protected]/50-budget.conf"))
+	if err != nil || !strings.Contains(string(data), "MemoryMax=16G\n") {
+		t.Fatalf("persistent shared ceiling: %q, %v", data, err)
+	}
+}
+
+func TestApplyScopes(t *testing.T) {
+	c := testConfig()
+	c.Profiles["compute"] = profile{"400%", "2G", "4G", "1G"}
+	c.Profiles["bulk"] = profile{"200%", "1G", "2G", "1G"}
+	c.Profiles["cache"] = profile{"100%", "1G", "3G", "0"}
+	f := newFake()
+	f.units = []unit{
+		{Name: "budget-c123abc.scope", Description: "compute workload", Active: "active"},
+		{Name: "budget-b123abc.scope", Description: "bulk workload", Active: "active"},
+		{Name: "budget-c456abc.scope", Description: "cache workload", Active: "active"},
+		{Name: "budget-s123abc.scope", Description: "search workload", Active: "active"},
+		{Name: "budget-c123def.scope", Description: "compute workload", Active: "inactive"},
+	}
+	a := application{c, homePaths(t.TempDir()), f, func(...string) error { return nil }, io.Discard}
+	if err := a.apply(); err != nil {
+		t.Fatal(err)
+	}
+	for name, want := range map[string]uint64{
+		"budget-c123abc.scope": 4294967296, "budget-b123abc.scope": 2147483648,
+		"budget-c456abc.scope": 3221225472,
+	} {
+		if got := values(f.sets[name])["MemoryMax"].Value(); got != want {
+			t.Errorf("%s ceiling: got %v, want %d", name, got, want)
+		}
+	}
+	if len(f.sets) != 3 {
+		t.Fatalf("unexpected scope updates: %v", keys(f.sets))
+	}
+}
+
+func TestApplyInvalid(t *testing.T) {
+	for name, change := range map[string]func(*config){
+		"resource": func(c *config) {
+			p := c.Profiles["default"]
+			p.High = "invalid"
+			c.Profiles["default"] = p
+		},
+		"command": func(c *config) { c.Commands["java"] = "missing" },
+		"prefix":  func(c *config) { c.Prefixes["default"] = []string{"snap-firefox"} },
+	} {
+		for _, active := range []bool{false, true} {
+			t.Run(fmt.Sprintf("%s/active=%t", name, active), func(t *testing.T) {
+				c := testConfig()
+				change(&c)
+				f := newFake()
+				if active {
+					f.units = []unit{{Name: "app-browser-123.scope", Active: "active"}}
+				}
+				var events []string
+				f.events = &events
+				p := homePaths(t.TempDir())
+				put(t, filepath.Join(p.units, "user.conf"), "keep")
+				a := application{c, p, f, func(args ...string) error {
+					events = append(events, strings.Join(args, " "))
+					return nil
+				}, io.Discard}
+				if err := a.apply(); err == nil {
+					t.Fatal("invalid configuration accepted")
+				}
+				if len(events) != 0 {
+					t.Fatal("invalid configuration changed units:", events)
+				}
+				entries, err := os.ReadDir(p.units)
+				if err != nil || len(entries) != 1 || entries[0].Name() != "user.conf" {
+					t.Fatalf("invalid configuration changed files: %v, %v", entries, err)
+				}
+				if _, err := os.Stat(p.state); !os.IsNotExist(err) {
+					t.Fatalf("invalid configuration created state: %v", err)
+				}
+			})
+		}
+	}
+}
+
 func TestSharedOverride(t *testing.T) {
 	f := newFake()
 	f.units = []unit{

internal/budget/config.go

@@ -72,7 +72,10 @@ func keys[V any](m map[string]V) []string {
 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) {
@@ -112,7 +115,11 @@ func (c config) properties(name string) ([]property, error) {
 	}
 	quota := uint64(math.MaxUint64)
 	if p.CPU != "" {
-		quota, err = strconv.ParseUint(strings.TrimSuffix(p.CPU, "%"), 10, 64)
+		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)
 		}

internal/budget/config_test.go

@@ -157,12 +157,24 @@ func TestLoadConfig(t *testing.T) {
 	data := `python = "compute"
 terminals = ["app-terminal-", "app-org.example.Console"]
 
+[profiles.default]
+cpu = "400%"
+high = "6G"
+memory = "8G"
+swap = "2G"
+
 [profiles.compute]
 cpu = "800%"
 high = "12G"
 memory = "16G"
 swap = "2G"
 
+[profiles.unrestricted]
+cpu = ""
+high = "infinity"
+memory = "infinity"
+swap = "infinity"
+
 [[devices]]
 path = '/dev/mock'
 [devices.limits.compute]
@@ -190,7 +202,9 @@ java = "compute"
 		t.Fatal(err)
 	}
 	want := testConfig()
-	want.Profiles = map[string]profile{"compute": want.Profiles["compute"]}
+	want.Profiles = map[string]profile{
+		"compute": want.Profiles["compute"], "default": want.Profiles["default"], "unrestricted": want.Profiles["unrestricted"],
+	}
 	want.Devices[0].Limits = map[string]ioLimit{"compute": want.Devices[0].Limits["compute"]}
 	want.Apps = map[string][]string{"compute": want.Apps["compute"]}
 	want.Units["unrestricted"] = []string{"containerd.service"}

internal/budget/files.go

@@ -37,7 +37,7 @@ func writeFiles(p paths, files map[string]string) error {
 	}
 	for _, relative := range previous {
 		if _, keep := files[relative]; !keep {
-			if err := os.Remove(filepath.Join(p.units, relative)); err != nil {
+			if err := os.Remove(filepath.Join(p.units, relative)); err != nil && !os.IsNotExist(err) {
 				return err
 			}
 		}

internal/budget/files_test.go

@@ -46,3 +46,40 @@ func TestWriteFiles(t *testing.T) {
 		t.Fatal("second apply failed:", err)
 	}
 }
+
+func TestFilesRetry(t *testing.T) {
+	p := homePaths(t.TempDir())
+	old := "old-a.scope.d/50-budget.conf"
+	put(t, filepath.Join(p.units, old), "old")
+	put(t, filepath.Join(p.state, "files.json"), `["old-a.scope.d/50-budget.conf", "old-b.scope.d/50-budget.conf"]`)
+	blocker := filepath.Join(p.units, "new.scope.d")
+	put(t, blocker, "blocked")
+	files := map[string]string{"new.scope.d/50-budget.conf": "[Scope]\nMemoryMax=8G\n"}
+	if err := writeFiles(p, files); err == nil {
+		t.Fatal("write failure ignored")
+	}
+	if _, err := os.Stat(filepath.Join(p.units, old)); !os.IsNotExist(err) {
+		t.Fatalf("cleanup did not finish before the write failure: %v", err)
+	}
+	if err := os.Remove(blocker); err != nil {
+		t.Fatal(err)
+	}
+	if err := writeFiles(p, files); err != nil {
+		t.Fatal("retry failed:", err)
+	}
+	data, err := os.ReadFile(filepath.Join(p.units, "new.scope.d/50-budget.conf"))
+	if err != nil || string(data) != files["new.scope.d/50-budget.conf"] {
+		t.Fatalf("replacement file: %q, %v", data, err)
+	}
+	data, err = os.ReadFile(filepath.Join(p.state, "files.json"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	var manifest []string
+	if err := json.Unmarshal(data, &manifest); err != nil {
+		t.Fatal(err)
+	}
+	if !reflect.DeepEqual(manifest, keys(files)) {
+		t.Fatal("manifest was not updated:", manifest)
+	}
+}

internal/budget/units.go

@@ -48,6 +48,11 @@ func (c config) unitProfile(unit string) string {
 			return name
 		}
 	}
+	for _, app := range keys(c.Shared) {
+		if appMatches(app, unit) {
+			return c.Shared[app]
+		}
+	}
 	for _, name := range keys(c.Apps) {
 		for _, app := range c.Apps[name] {
 			if appMatches(app, unit) {
@@ -57,10 +62,7 @@ func (c config) unitProfile(unit string) string {
 	}
 	for _, name := range keys(c.Prefixes) {
 		for _, prefix := range c.Prefixes[name] {
-			parts := strings.Split(prefix, ".")
+			if prefixMatches(prefix, unit) {
-			stem := strings.Join(parts[:len(parts)-2], ".")
-			kind := "." + parts[len(parts)-2]
-			if strings.HasPrefix(unit, stem) && strings.HasSuffix(unit, kind) {
 				return name
 			}
 		}
@@ -77,16 +79,51 @@ func (c config) unitProfile(unit string) string {
 	return ""
 }
 
+func (c config) activeProfile(u unit) string {
+	if name := c.unitProfile(u.Name); name != "" {
+		return name
+	}
+	if strings.HasPrefix(u.Name, "budget-") && strings.HasSuffix(u.Name, ".scope") {
+		if name, ok := strings.CutSuffix(u.Description, " workload"); ok && name != "search" {
+			return name
+		}
+	}
+	return ""
+}
+
+func prefixMatches(directory, unit string) bool {
+	for _, kind := range []string{".service", ".scope", ".slice"} {
+		if stem, ok := strings.CutSuffix(directory, kind+".d"); ok {
+			return strings.HasPrefix(unit, stem) && strings.HasSuffix(unit, kind)
+		}
+	}
+	return false
+}
+
+func dropSection(directory string) (string, error) {
+	for _, kind := range []struct{ suffix, section string }{
+		{".service.d", "Service"}, {".scope.d", "Scope"}, {".slice.d", "Slice"},
+	} {
+		if stem, ok := strings.CutSuffix(directory, kind.suffix); ok && stem != "" {
+			return kind.section, nil
+		}
+	}
+	return "", fmt.Errorf("invalid drop-in directory %q", directory)
+}
+
 func (c config) render() (map[string]string, error) {
+	if err := c.validate(); err != nil {
+		return nil, err
+	}
 	files := map[string]string{}
 	add := func(directory, name, extra string) error {
 		values, err := c.assignments(name)
 		if err != nil {
 			return err
 		}
-		section := "Service"
+		section, err := dropSection(directory)
-		if strings.Contains(directory, ".scope.") {
+		if err != nil {
-			section = "Scope"
+			return err
 		}
 		files[directory+"/50-budget.conf"] = "[" + section + "]\n" + strings.Join(values, "\n") + "\n" + extra
 		return nil

internal/budget/units_test.go

@@ -98,3 +98,38 @@ func TestNoSharedApps(t *testing.T) {
 		t.Fatal(got)
 	}
 }
+
+func TestSliceDropIns(t *testing.T) {
+	c := testConfig()
+	c.Units["compute"] = []string{"custom.slice"}
+	c.Prefixes["bulk"] = []string{"custom-.slice.d"}
+	files, err := c.render()
+	if err != nil {
+		t.Fatal(err)
+	}
+	for file, limit := range map[string]string{
+		"custom.slice.d/50-budget.conf": "16G", "custom-.slice.d/50-budget.conf": "8G",
+	} {
+		if text := files[file]; !strings.HasPrefix(text, "[Slice]\n") || !strings.Contains(text, "MemoryMax="+limit+"\n") {
+			t.Fatalf("%s: %q", file, text)
+		}
+	}
+	if got := c.unitProfile("custom-work.slice"); got != "bulk" {
+		t.Fatalf("slice prefix profile: %q", got)
+	}
+}
+
+func TestInvalidPrefixes(t *testing.T) {
+	for _, prefix := range []string{"snap-firefox", "snap.firefox", "snap.firefox.scope", "snap.firefox.timer.d"} {
+		t.Run(prefix, func(t *testing.T) {
+			c := testConfig()
+			c.Prefixes["default"] = []string{prefix}
+			if _, err := c.render(); err == nil {
+				t.Fatal("invalid prefix accepted")
+			}
+			if got := c.unitProfile("session-ui.service"); got != "" {
+				t.Fatalf("invalid prefix matched an unrelated unit: %q", got)
+			}
+		})
+	}
+}

internal/budget/validate.go

@@ -0,0 +1,61 @@
+package budget
+
+import "fmt"
+
+func (c config) validate() error {
+	for _, name := range keys(c.Profiles) {
+		if _, err := c.properties(name); err != nil {
+			return fmt.Errorf("profiles.%s: %w", name, err)
+		}
+	}
+	check := func(field, name string) error {
+		if _, err := c.getProfile(name); err != nil {
+			return fmt.Errorf("%s: %w", field, err)
+		}
+		return nil
+	}
+	for _, field := range []struct {
+		name   string
+		groups map[string][]string
+	}{
+		{"apps", c.Apps}, {"units", c.Units}, {"prefixes", c.Prefixes},
+	} {
+		for _, name := range keys(field.groups) {
+			if err := check(field.name+"."+name, name); err != nil {
+				return err
+			}
+		}
+	}
+	for _, field := range []struct {
+		name string
+		refs map[string]string
+	}{
+		{"commands", c.Commands}, {"shared", c.Shared},
+	} {
+		for _, key := range keys(field.refs) {
+			if err := check(field.name+"."+key, field.refs[key]); err != nil {
+				return err
+			}
+		}
+	}
+	if c.Python != "" {
+		if err := check("python", c.Python); err != nil {
+			return err
+		}
+	}
+	for i, device := range c.Devices {
+		for _, name := range keys(device.Limits) {
+			if err := check(fmt.Sprintf("devices[%d].limits.%s", i, name), name); err != nil {
+				return err
+			}
+		}
+	}
+	for _, name := range keys(c.Prefixes) {
+		for _, directory := range c.Prefixes[name] {
+			if _, err := dropSection(directory); err != nil {
+				return fmt.Errorf("prefixes.%s: %w", name, err)
+			}
+		}
+	}
+	return nil
+}

internal/budget/validate_test.go

@@ -0,0 +1,74 @@
+package budget
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/BurntSushi/toml"
+)
+
+func TestConfigInvalid(t *testing.T) {
+	for _, tc := range []struct {
+		name   string
+		field  string
+		change func(*config)
+	}{
+		{"cpu", "profiles.compute", func(c *config) {
+			p := c.Profiles["compute"]
+			p.CPU = "invalid%"
+			c.Profiles["compute"] = p
+		}},
+		{"cpu-suffix", "profiles.compute", func(c *config) {
+			p := c.Profiles["compute"]
+			p.CPU = "800"
+			c.Profiles["compute"] = p
+		}},
+		{"high", "profiles.compute", func(c *config) {
+			p := c.Profiles["compute"]
+			p.High = "invalid"
+			c.Profiles["compute"] = p
+		}},
+		{"memory", "profiles.compute", func(c *config) {
+			p := c.Profiles["compute"]
+			p.Memory = "invalid"
+			c.Profiles["compute"] = p
+		}},
+		{"swap", "profiles.compute", func(c *config) {
+			p := c.Profiles["compute"]
+			p.Swap = "invalid"
+			c.Profiles["compute"] = p
+		}},
+		{"io", "profiles.compute", func(c *config) {
+			c.Devices[0].Limits["compute"] = ioLimit{"invalid", "500M", "20000", "20000"}
+		}},
+		{"command", "commands.java", func(c *config) { c.Commands["java"] = "missing" }},
+		{"python", "python", func(c *config) { c.Python = "missing" }},
+		{"shared", "shared.editor", func(c *config) { c.Shared["editor"] = "missing" }},
+		{"apps", "apps.missing", func(c *config) { c.Apps["missing"] = []string{"editor"} }},
+		{"units", "units.missing", func(c *config) { c.Units["missing"] = []string{"worker.service"} }},
+		{"prefixes", "prefixes.missing", func(c *config) { c.Prefixes["missing"] = []string{"app-editor-.scope.d"} }},
+		{"device", "devices[0].limits", func(c *config) {
+			c.Devices[0].Limits["missing"] = ioLimit{"100M", "100M", "4000", "4000"}
+		}},
+		{"prefix", "prefixes.default", func(c *config) { c.Prefixes["default"] = []string{"snap-firefox"} }},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			c := testConfig()
+			tc.change(&c)
+			var data bytes.Buffer
+			if err := toml.NewEncoder(&data).Encode(c); err != nil {
+				t.Fatal(err)
+			}
+			file := filepath.Join(t.TempDir(), "profiles.toml")
+			if err := os.WriteFile(file, data.Bytes(), 0600); err != nil {
+				t.Fatal(err)
+			}
+			if _, err := loadConfig(file); err == nil || !strings.Contains(err.Error(), tc.field) {
+				t.Fatalf("expected an error naming %s, got %v", tc.field, err)
+			}
+		})
+	}
+}