feat: add workloads to the plan

internals/overlord/planstate/manager.go

@@ -253,6 +253,8 @@ func (m *PlanManager) updatePlanLayers(layers []*plan.Layer) (*plan.Plan, error)
 	if err != nil {
 		return nil, err
 	}
+	// Validate workloads & services here?
+	// We can make a switch in personality for KernOS/Pebble
 	m.plan = p
 	return p, nil
 }

internals/overlord/servstate/handlers.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"maps"
 	"os"
 	"os/exec"
 	"os/user"
@@ -95,6 +96,7 @@ type serviceData struct {
 	manager      *ServiceManager
 	state        serviceState
 	config       *plan.Service
+	workload     *Workload
 	logs         *servicelog.RingBuffer
 	started      chan error
 	stopped      chan error
@@ -115,13 +117,21 @@ func (m *ServiceManager) doStart(task *state.Task, tomb *tomb.Tomb) error {
 	}
 
 	currentPlan := m.getPlan()
-	config, ok := currentPlan.Services[request.Name]
-	if !ok {
+	config, configFound := currentPlan.Services[request.Name]
+	if !configFound {
 		return fmt.Errorf("cannot find service %q in plan", request.Name)
 	}
 
+	var workload *Workload
+	if config.Workload != "" {
+		ws := currentPlan.Sections[WorkloadsField].(*WorkloadsSection)
+		if workload = ws.Entries[config.Workload]; workload == nil {
+			return fmt.Errorf("cannot find workload %q for service %q in plan", config.Workload, request.Name)
+		}
+	}
+
 	// Create the service object (or reuse the existing one by name).
-	service, taskLog := m.serviceForStart(config)
+	service, taskLog := m.serviceForStart(config, workload)
 	if taskLog != "" {
 		addTaskLog(task, taskLog)
 	}
@@ -167,7 +177,7 @@ func (m *ServiceManager) doStart(task *state.Task, tomb *tomb.Tomb) error {
 // and is running.
 //
 // It also returns a message to add to the task's log, or empty string if none.
-func (m *ServiceManager) serviceForStart(config *plan.Service) (service *serviceData, taskLog string) {
+func (m *ServiceManager) serviceForStart(config *plan.Service, workload *Workload) (service *serviceData, taskLog string) {
 	m.servicesLock.Lock()
 	defer m.servicesLock.Unlock()
 
@@ -177,17 +187,23 @@ func (m *ServiceManager) serviceForStart(config *plan.Service) (service *service
 		service = &serviceData{
 			manager: m,
 			state:   stateInitial,
-			config:  config.Copy(),
 			logs:    servicelog.NewRingBuffer(maxLogBytes),
 			started: make(chan error, 1),
 			stopped: make(chan error, 2), // enough for killTimeElapsed to send, and exit if it happens after
 		}
+		service.config = config.Copy()
+		if workload != nil {
+			service.workload = workload.copy()
+		}
 		m.services[config.Name] = service
 		return service, ""
 	}
 
 	// Ensure config is up-to-date from the plan whenever the user starts a service.
 	service.config = config.Copy()
+	if workload != nil {
+		service.workload = workload.copy()
+	}
 
 	switch service.state {
 	case stateInitial, stateStarting, stateRunning:
@@ -338,17 +354,31 @@ func (s *serviceData) startInternal() error {
 	s.cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
 
 	// Copy environment to avoid updating original.
-	environment := make(map[string]string)
-	for k, v := range s.config.Environment {
-		environment[k] = v
+	var environment map[string]string
+	if s.workload != nil && len(s.workload.Environment) > 0 {
+		environment = maps.Clone(s.workload.Environment)
+	} else if len(s.config.Environment) > 0 {
+		environment = maps.Clone(s.config.Environment)
+	} else {
+		environment = make(map[string]string)
 	}
 
 	s.cmd.Dir = s.config.WorkingDir
 
 	// Start as another user if specified in plan.
-	uid, gid, err := osutil.NormalizeUidGid(s.config.UserID, s.config.GroupID, s.config.User, s.config.Group)
-	if err != nil {
-		return err
+	var uid, gid *int
+	if s.config.UserID != nil || s.config.GroupID != nil || s.config.User != "" || s.config.Group != "" {
+		// User/group config from the service takes precedence if any of them are set
+		uid, gid, err = osutil.NormalizeUidGid(s.config.UserID, s.config.GroupID, s.config.User, s.config.Group)
+		if err != nil {
+			return err
+		}
+	} else if s.workload != nil {
+		// Take user/group config from workload
+		uid, gid, err = osutil.NormalizeUidGid(s.workload.UserID, s.workload.GroupID, s.workload.User, s.workload.Group)
+		if err != nil {
+			return err
+		}
 	}
 	if uid != nil && gid != nil {
 		isCurrent, err := osutil.IsCurrent(*uid, *gid)

internals/overlord/servstate/manager.go

@@ -264,7 +264,17 @@ func (m *ServiceManager) Replan() ([][]string, [][]string, error) {
 			if config.Equal(s.config) {
 				continue
 			}
-			s.config = config.Copy() // update service config from plan
+			// Update service config and workload from plan
+			s.config = config.Copy()
+			if s.config.Workload != "" {
+				ws, ok := currentPlan.Sections[WorkloadsField].(*WorkloadsSection)
+				if !ok {
+					return nil, nil, fmt.Errorf("internal error: invalid section type %T", ws)
+				}
+				s.workload = ws.Entries[s.config.Workload].copy()
+			} else {
+				s.workload = nil
+			}
 		}
 		needsRestart[name] = true
 		stop = append(stop, name)

internals/overlord/servstate/manager_test.go

@@ -175,6 +175,7 @@ func (s *S) TearDownTest(c *C) {
 			s.stopRunningServices(c)
 		}
 	}
+
 	// General test cleanup
 	s.BaseTest.TearDownTest(c)
 
@@ -1797,12 +1798,15 @@ func (s *S) planAddLayer(c *C, layerYAML string) {
 	layers := append(s.plan.Layers, layer)
 	combined, err := plan.CombineLayers(layers...)
 	c.Assert(err, IsNil)
+	c.Assert(combined.Validate(), IsNil)
 	s.plan = &plan.Plan{
 		Layers:     layers,
 		Services:   combined.Services,
 		Checks:     combined.Checks,
 		LogTargets: combined.LogTargets,
+		Sections:   combined.Sections,
 	}
+	c.Assert(s.plan.Validate(), IsNil)
 }
 
 // Make sure services are all stopped before the next test starts.

internals/overlord/servstate/workloads.go

@@ -0,0 +1,217 @@
+// Copyright (c) 2025 Canonical Ltd
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package servstate
+
+import (
+	"bytes"
+	"errors"
+	"fmt"
+	"maps"
+
+	"gopkg.in/yaml.v3"
+
+	"github.com/canonical/pebble/internals/osutil"
+	"github.com/canonical/pebble/internals/plan"
+)
+
+var _ plan.SectionExtension = (*WorkloadsSectionExtension)(nil)
+
+type WorkloadsSectionExtension struct{}
+
+func (ext *WorkloadsSectionExtension) CombineSections(sections ...plan.Section) (plan.Section, error) {
+	ws := &WorkloadsSection{}
+	for _, section := range sections {
+		layer, ok := section.(*WorkloadsSection)
+		if !ok {
+			return nil, fmt.Errorf("internal error: invalid section type %T", layer)
+		}
+		if err := ws.combine(layer); err != nil {
+			return nil, err
+		}
+	}
+	return ws, nil
+}
+
+func (ext *WorkloadsSectionExtension) ParseSection(data yaml.Node) (plan.Section, error) {
+	ws := &WorkloadsSection{}
+	// The following issue prevents us from using the yaml.Node decoder
+	// with KnownFields = true behavior. Once one of the proposals get
+	// merged, we can remove the intermediate Marshal step.
+	// https://github.com/go-yaml/yaml/issues/460
+	if len(data.Content) != 0 {
+		yml, err := yaml.Marshal(data)
+		if err != nil {
+			return nil, fmt.Errorf(`internal error: cannot marshal "workloads" section: %w`, err)
+		}
+		dec := yaml.NewDecoder(bytes.NewReader(yml))
+		dec.KnownFields(true)
+		if err = dec.Decode(ws); err != nil {
+			return nil, &plan.FormatError{
+				Message: fmt.Sprintf(`cannot parse the "workloads" section: %v`, err),
+			}
+		}
+	}
+	for name, workload := range ws.Entries {
+		if workload != nil {
+			workload.Name = name
+		}
+	}
+	return ws, nil
+}
+
+func (ext *WorkloadsSectionExtension) ValidatePlan(p *plan.Plan) error {
+	ws, ok := p.Sections[WorkloadsField].(*WorkloadsSection)
+	if !ok {
+		return fmt.Errorf("internal error: invalid section type %T", ws)
+	}
+	for name, service := range p.Services {
+		if service.Workload == "" {
+			continue
+		}
+		if _, ok := ws.Entries[service.Workload]; !ok {
+			return &plan.FormatError{
+				Message: fmt.Sprintf(`plan service %q workload not defined: %q`, name, service.Workload),
+			}
+		}
+	}
+	for name, workload := range ws.Entries {
+		if _, _, err := osutil.NormalizeUidGid(workload.UserID, workload.GroupID, workload.User, workload.Group); err != nil {
+			return &plan.FormatError{
+				Message: fmt.Sprintf(`plan workload %q %v`, err, name),
+			}
+		}
+	}
+	return nil
+}
+
+const WorkloadsField = "workloads"
+
+var _ plan.Section = (*WorkloadsSection)(nil)
+
+type WorkloadsSection struct {
+	Entries map[string]*Workload `yaml:",inline"`
+}
+
+func (ws *WorkloadsSection) IsZero() bool {
+	return len(ws.Entries) == 0
+}
+
+func (ws *WorkloadsSection) Validate() error {
+	for name, workload := range ws.Entries {
+		if workload == nil {
+			return &plan.FormatError{
+				Message: fmt.Sprintf("workload %q cannot have a null value", name),
+			}
+		}
+		if err := workload.validate(); err != nil {
+			return &plan.FormatError{
+				Message: fmt.Sprintf("workload %q %v", name, err),
+			}
+		}
+	}
+	return nil
+}
+
+func (ws *WorkloadsSection) combine(other *WorkloadsSection) error {
+	for name, workload := range other.Entries {
+		if ws.Entries == nil {
+			ws.Entries = make(map[string]*Workload, len(other.Entries))
+		}
+		switch workload.Override {
+		case plan.MergeOverride:
+			if current, ok := ws.Entries[name]; ok {
+				copied := current.copy()
+				copied.merge(workload)
+				ws.Entries[name] = copied
+				break
+			}
+			fallthrough
+		case plan.ReplaceOverride:
+			ws.Entries[name] = workload.copy()
+		case plan.UnknownOverride:
+			return &plan.FormatError{
+				Message: fmt.Sprintf(`workload %q must define an "override" policy`, name),
+			}
+		default:
+			return &plan.FormatError{
+				Message: fmt.Sprintf(`workload %q has an invalid "override" policy: %q`, name, workload.Override),
+			}
+		}
+	}
+	return nil
+}
+
+type Workload struct {
+	// Basic details
+	Name     string        `yaml:"-"`
+	Override plan.Override `yaml:"override,omitempty"`
+
+	// Options for command execution
+	Environment map[string]string `yaml:"environment,omitempty"`
+	UserID      *int              `yaml:"user-id,omitempty"`
+	User        string            `yaml:"user,omitempty"`
+	GroupID     *int              `yaml:"group-id,omitempty"`
+	Group       string            `yaml:"group,omitempty"`
+}
+
+func (w *Workload) validate() error {
+	if w.Name == "" {
+		return errors.New("cannot have an empty name")
+	}
+	// Value of Override is checked in the (*WorkloadSection).combine() method
+	return nil
+}
+
+func (w *Workload) copy() *Workload {
+	copied := *w
+	copied.Environment = maps.Clone(w.Environment)
+	copied.UserID = copyPtr(w.UserID)
+	copied.GroupID = copyPtr(w.GroupID)
+	return &copied
+}
+
+func (w *Workload) merge(other *Workload) {
+	if len(other.Environment) > 0 {
+		w.Environment = makeMapIfNil(w.Environment)
+		maps.Copy(w.Environment, other.Environment)
+	}
+	if other.UserID != nil {
+		w.UserID = copyPtr(other.UserID)
+	}
+	if other.User != "" {
+		w.User = other.User
+	}
+	if other.GroupID != nil {
+		w.GroupID = copyPtr(other.GroupID)
+	}
+	if other.Group != "" {
+		w.Group = other.Group
+	}
+}
+
+func copyPtr[T any](p *T) *T {
+	if p == nil {
+		return nil
+	}
+	copied := *p
+	return &copied
+}
+
+func makeMapIfNil[K comparable, V any](m map[K]V) map[K]V {
+	if m == nil {
+		m = make(map[K]V)
+	}
+	return m
+}