Merge pull request #26 from kolyshkin/add-pid

Implement AddPid for cgroup managers

Author: kolyshkin

cgroups.go

@@ -29,6 +29,11 @@ type Manager interface {
 	// can be used to merely create a cgroup.
 	Apply(pid int) error
 
+	// AddPid adds a process with a given pid to an existing cgroup.
+	// The subcgroup argument is either empty, or a path relative to
+	// a cgroup under under the manager's cgroup.
+	AddPid(subcgroup string, pid int) error
+
 	// GetPids returns the PIDs of all processes inside the cgroup.
 	GetPids() ([]int, error)
 

fs/fs.go

@@ -4,6 +4,8 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"path"
+	"strings"
 	"sync"
 
 	"golang.org/x/sys/unix"
@@ -139,6 +141,33 @@ func (m *Manager) Apply(pid int) (retErr error) {
 	return retErr
 }
 
+// AddPid adds a process with a given pid to an existing cgroup.
+// The subcgroup argument is either empty, or a path relative to
+// a cgroup under under the manager's cgroup.
+func (m *Manager) AddPid(subcgroup string, pid int) (retErr error) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	c := m.cgroups
+
+	for _, dir := range m.paths {
+		path := path.Join(dir, subcgroup)
+		if !strings.HasPrefix(path, dir) {
+			return fmt.Errorf("bad sub cgroup path: %s", subcgroup)
+		}
+
+		if err := cgroups.WriteCgroupProc(path, pid); err != nil {
+			if isIgnorableError(c.Rootless, err) && c.Path == "" {
+				retErr = cgroups.ErrRootless
+				continue
+			}
+			return err
+		}
+	}
+
+	return retErr
+}
+
 func (m *Manager) Destroy() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()

fs2/fs2.go

@@ -4,6 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"path/filepath"
 	"strings"
 
 	"github.com/opencontainers/cgroups"
@@ -83,6 +84,18 @@ func (m *Manager) Apply(pid int) error {
 	return nil
 }
 
+// AddPid adds a process with a given pid to an existing cgroup.
+// The subcgroup argument is either empty, or a path relative to
+// a cgroup under under the manager's cgroup.
+func (m *Manager) AddPid(subcgroup string, pid int) error {
+	path := filepath.Join(m.dirPath, subcgroup)
+	if !strings.HasPrefix(path, m.dirPath) {
+		return fmt.Errorf("bad sub cgroup path: %s", subcgroup)
+	}
+
+	return cgroups.WriteCgroupProc(path, pid)
+}
+
 func (m *Manager) GetPids() ([]int, error) {
 	return cgroups.GetPids(m.dirPath)
 }

go.mod

@@ -4,7 +4,7 @@ go 1.23.0
 
 require (
 	github.com/cilium/ebpf v0.17.3
-	github.com/coreos/go-systemd/v22 v22.5.0
+	github.com/coreos/go-systemd/v22 v22.6.0
 	github.com/cyphar/filepath-securejoin v0.4.1
 	github.com/godbus/dbus/v5 v5.1.0
 	github.com/moby/sys/mountinfo v0.7.2

go.sum

@@ -1,15 +1,14 @@
 github.com/cilium/ebpf v0.17.3 h1:FnP4r16PWYSE4ux6zN+//jMcW4nMVRvuTLVTvCjyyjg=
 github.com/cilium/ebpf v0.17.3/go.mod h1:G5EDHij8yiLzaqn0WjyfJHvRa+3aDlReIaLVRMvOyJk=
-github.com/coreos/go-systemd/v22 v22.5.0 h1:RrqgGjYQKalulkV8NGVIfkXQf6YYmOyiJKk8iXXhfZs=
-github.com/coreos/go-systemd/v22 v22.5.0/go.mod h1:Y58oyj3AT4RCenI/lSvhwexgC+NSVTIJ3seZv2GcEnc=
+github.com/coreos/go-systemd/v22 v22.6.0 h1:aGVa/v8B7hpb0TKl0MWoAavPDmHvobFe5R5zn0bCJWo=
+github.com/coreos/go-systemd/v22 v22.6.0/go.mod h1:iG+pp635Fo7ZmV/j14KUcmEyWF+0X7Lua8rrTWzYgWU=
 github.com/cyphar/filepath-securejoin v0.4.1 h1:JyxxyPEaktOD+GAnqIqTf9A8tHyAG22rowi7HkoSU1s=
 github.com/cyphar/filepath-securejoin v0.4.1/go.mod h1:Sdj7gXlvMcPZsbhwhQ33GguGLDGQL7h7bg04C/+u9jI=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/go-quicktest/qt v1.101.0 h1:O1K29Txy5P2OK0dGo59b7b0LR6wKfIhttaAhHUyn7eI=
 github.com/go-quicktest/qt v1.101.0/go.mod h1:14Bz/f7NwaXPtdYEgzsx46kqSxVwTbzVZsDC26tQJow=
-github.com/godbus/dbus/v5 v5.0.4/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/godbus/dbus/v5 v5.1.0 h1:4KLkAxT3aOY8Li4FRJe/KvhoNFFxo0m6fNuFUO8QJUk=
 github.com/godbus/dbus/v5 v5.1.0/go.mod h1:xhWf0FNVPg57R7Z0UbKHbJfkEywrmjJnf7w5xrFpKfA=
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=

systemd/common.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math"
 	"os"
+	"path"
 	"strconv"
 	"strings"
 	"sync"
@@ -208,6 +209,20 @@ func stopUnit(cm *dbusConnManager, unitName string) error {
 	return nil
 }
 
+func addPid(cm *dbusConnManager, unitName, subcgroup string, pid int) error {
+	absSubcgroup := subcgroup
+	if !path.IsAbs(absSubcgroup) {
+		absSubcgroup = "/" + subcgroup
+	}
+	if absSubcgroup != path.Clean(absSubcgroup) {
+		return fmt.Errorf("bad sub cgroup path: %s", subcgroup)
+	}
+
+	return cm.retryOnDisconnect(func(c *systemdDbus.Conn) error {
+		return c.AttachProcessesToUnit(context.TODO(), unitName, absSubcgroup, []uint32{uint32(pid)})
+	})
+}
+
 func resetFailedUnit(cm *dbusConnManager, name string) error {
 	return cm.retryOnDisconnect(func(c *systemdDbus.Conn) error {
 		return c.ResetFailedUnitContext(context.TODO(), name)

systemd/v1.go

@@ -215,6 +215,25 @@ func (m *LegacyManager) Apply(pid int) error {
 	return nil
 }
 
+// AddPid adds a process with a given pid to an existing cgroup.
+// The subcgroup argument is either empty, or a path relative to
+// a cgroup under under the manager's cgroup.
+func (m *LegacyManager) AddPid(subcgroup string, pid int) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	if err := addPid(m.dbus, getUnitName(m.cgroups), subcgroup, pid); err != nil {
+		return err
+	}
+
+	// Since systemd only joins controllers it knows, use cgroupfs for the rest.
+	fsMgr, err := fs.NewManager(m.cgroups, m.paths)
+	if err != nil {
+		return err
+	}
+	return fsMgr.AddPid(subcgroup, pid)
+}
+
 func (m *LegacyManager) Destroy() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()

systemd/v2.go

@@ -383,6 +383,16 @@ func cgroupFilesToChown() ([]string, error) {
 	return filesToChown, nil
 }
 
+// AddPid adds a process with a given pid to an existing cgroup.
+// The subcgroup argument is either empty, or a path relative to
+// a cgroup under under the manager's cgroup.
+func (m *UnifiedManager) AddPid(subcgroup string, pid int) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	return addPid(m.dbus, getUnitName(m.cgroups), subcgroup, pid)
+}
+
 func (m *UnifiedManager) Destroy() error {
 	m.mu.Lock()
 	defer m.mu.Unlock()