Merge pull request #203 from sameo/topic/subsystem-filter

cgroup: Optionally add process and task to a subsystems subset

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

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 {

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