cgroup: Optionally add process and task to a subsystems subset

sameo · Samuel Ortiz · commit 80a7821536fc · 2021-10-07T08:36:11.000+02:00
In some cases, one may want to add processes and task only to certain
subsystems in a cgroup. For example, when a cgroup is created to only
limit a given subsystem usage, but leaves all other subsystems
unconstrained, there may be a need to move a process to only the
contrained cgroup subsystem while leaving it in another cgroup for
other subsystems.

This commit makes the Control interface process and task addition
functions variadic, in order to pass an optional list of subsystems that
behave as s subsystems filter.

Signed-off-by: Samuel Ortiz &lt;samuel.e.ortiz@protonmail.com&gt;
diff --git a/cgroup.go b/cgroup.go
@@ -150,22 +150,50 @@ func (c *cgroup) Subsystems() []Subsystem {
 	return c.subsystems
 }
 
-// Add moves the provided process into the new cgroup
-func (c *cgroup) Add(process Process) error {
-	return c.add(process, cgroupProcs)
+func (c *cgroup) subsystemsFilter(subsystems ...Name) []Subsystem {
+	if len(subsystems) == 0 {
+		return c.subsystems
+	}
+
+	var filteredSubsystems = []Subsystem{}
+	for _, s := range c.subsystems {
+		for _, f := range subsystems {
+			if s.Name() == f {
+				filteredSubsystems = append(filteredSubsystems, s)
+				break
+			}
+		}
+	}
+
+	return filteredSubsystems
 }
 
-// AddProc moves the provided process id into the new cgroup
-func (c *cgroup) AddProc(pid uint64) error {
-	return c.add(Process{Pid: int(pid)}, cgroupProcs)
+// Add moves the provided process into the new cgroup.
+// Without additional arguments, the process is added to all the cgroup subsystems.
+// When giving Add a list of subsystem names, the process is only added to those
+// subsystems, provided that they are active in the targeted cgroup.
+func (c *cgroup) Add(process Process, subsystems ...Name) error {
+	return c.add(process, cgroupProcs, subsystems...)
 }
 
-// AddTask moves the provided tasks (threads) into the new cgroup
-func (c *cgroup) AddTask(process Process) error {
-	return c.add(process, cgroupTasks)
+// AddProc moves the provided process id into the new cgroup.
+// Without additional arguments, the process with the given id is added to all
+// the cgroup subsystems. When giving AddProc a list of subsystem names, the process
+// id is only added to those subsystems, provided that they are active in the targeted
+// cgroup.
+func (c *cgroup) AddProc(pid uint64, subsystems ...Name) error {
+	return c.add(Process{Pid: int(pid)}, cgroupProcs, subsystems...)
 }
 
-func (c *cgroup) add(process Process, pType procType) error {
+// AddTask moves the provided tasks (threads) into the new cgroup.
+// Without additional arguments, the task is added to all the cgroup subsystems.
+// When giving AddTask a list of subsystem names, the task is only added to those
+// subsystems, provided that they are active in the targeted cgroup.
+func (c *cgroup) AddTask(process Process, subsystems ...Name) error {
+	return c.add(process, cgroupTasks, subsystems...)
+}
+
+func (c *cgroup) add(process Process, pType procType, subsystems ...Name) error {
 	if process.Pid <= 0 {
 		return ErrInvalidPid
 	}
@@ -174,7 +202,7 @@ func (c *cgroup) add(process Process, pType procType) error {
 	if c.err != nil {
 		return c.err
 	}
-	for _, s := range pathers(c.subsystems) {
+	for _, s := range pathers(c.subsystemsFilter(subsystems...)) {
 		p, err := c.path(s.Name())
 		if err != nil {
 			return err
diff --git a/cgroup_test.go b/cgroup_test.go
@@ -101,6 +101,63 @@ func TestAdd(t *testing.T) {
 	}
 }
 
+func TestAddFilteredSubsystems(t *testing.T) {
+	mock, err := newMock()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer mock.delete()
+	control, err := New(mock.hierarchy, StaticPath("test"), &specs.LinuxResources{})
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	filteredSubsystems := []Name{"memory", "cpu"}
+	if err := control.Add(Process{Pid: 1234}, filteredSubsystems...); err != nil {
+		t.Error(err)
+		return
+	}
+
+	for _, s := range filteredSubsystems {
+		if err := checkPid(mock, filepath.Join(string(s), "test"), 1234); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+
+	if err := checkPid(mock, filepath.Join("devices", "test"), 1234); err == nil {
+		t.Error("Pid should not be added to the devices subsystem")
+		return
+	}
+
+	bogusSubsystems := append(filteredSubsystems, "bogus")
+	if err := control.Add(Process{Pid: 5678}, bogusSubsystems...); err != nil {
+		t.Error(err)
+		return
+	}
+
+	for _, s := range filteredSubsystems {
+		if err := checkPid(mock, filepath.Join(string(s), "test"), 5678); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+
+	nilSubsystems := []Name{}
+	if err := control.Add(Process{Pid: 9012}, nilSubsystems...); err != nil {
+		t.Error(err)
+		return
+	}
+
+	for _, s := range Subsystems() {
+		if err := checkPid(mock, filepath.Join(string(s), "test"), 9012); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+}
+
 func TestAddTask(t *testing.T) {
 	mock, err := newMock()
 	if err != nil {
@@ -124,6 +181,48 @@ func TestAddTask(t *testing.T) {
 	}
 }
 
+func TestAddTaskFilteredSubsystems(t *testing.T) {
+	mock, err := newMock()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer mock.delete()
+	control, err := New(mock.hierarchy, StaticPath("test"), &specs.LinuxResources{})
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	filteredSubsystems := []Name{"memory", "cpu"}
+	if err := control.AddTask(Process{Pid: 1234}, filteredSubsystems...); err != nil {
+		t.Error(err)
+		return
+	}
+	for _, s := range filteredSubsystems {
+		if err := checkTaskid(mock, filepath.Join(string(s), "test"), 1234); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+
+	if err := checkTaskid(mock, filepath.Join("devices", "test"), 1234); err == nil {
+		t.Error("Task should not be added to the devices subsystem")
+		return
+	}
+
+	bogusSubsystems := append(filteredSubsystems, "bogus")
+	if err := control.AddTask(Process{Pid: 5678}, bogusSubsystems...); err != nil {
+		t.Error(err)
+		return
+	}
+
+	for _, s := range filteredSubsystems {
+		if err := checkTaskid(mock, filepath.Join(string(s), "test"), 5678); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+}
+
 func TestListPids(t *testing.T) {
 	mock, err := newMock()
 	if err != nil {
diff --git a/control.go b/control.go
@@ -54,12 +54,22 @@ type Task = Process
 type Cgroup interface {
 	// New creates a new cgroup under the calling cgroup
 	New(string, *specs.LinuxResources) (Cgroup, error)
-	// Add adds a process to the cgroup (cgroup.procs)
-	Add(Process) error
-	// AddProc adds the process with the given id to the cgroup (cgroup.procs)
-	AddProc(pid uint64) error
-	// AddTask adds a process to the cgroup (tasks)
-	AddTask(Process) error
+	// Add adds a process to the cgroup (cgroup.procs). Without additional arguments,
+	// the process is added to all the cgroup subsystems. When giving Add a list of
+	// subsystem names, the process is only added to those subsystems, provided that
+	// they are active in the targeted cgroup.
+	Add(Process, ...Name) error
+	// AddProc adds the process with the given id to the cgroup (cgroup.procs).
+	// Without additional arguments, the process with the given id is added to all
+	// the cgroup subsystems. When giving AddProc a list of subsystem names, the process
+	// id is only added to those subsystems, provided that they are active in the targeted
+	// cgroup.
+	AddProc(uint64, ...Name) error
+	// AddTask adds a process to the cgroup (tasks). Without additional arguments, the
+	// task is added to all the cgroup subsystems. When giving AddTask a list of subsystem
+	// names, the task is only added to those subsystems, provided that they are active in
+	// the targeted cgroup.
+	AddTask(Process, ...Name) error
 	// Delete removes the cgroup as a whole
 	Delete() error
 	// MoveTo moves all the processes under the calling cgroup to the provided one
