Add functionality for retrieving all tasks of a cgroup

gpanouts · gpanouts · commit a31a0ff98523 · 2018-11-21T10:12:58.000+02:00
Signed-off-by: Georgia Panoutsakopoulou &lt;gpanoutsak@gmail.com&gt;
diff --git a/cgroup.go b/cgroup.go
@@ -321,6 +321,49 @@ func (c *cgroup) processes(subsystem Name, recursive bool) ([]Process, error) {
 	return processes, err
 }
 
+// Tasks returns the tasks running inside the cgroup along
+// with the subsystem used, pid, and path
+func (c *cgroup) Tasks(subsystem Name, recursive bool) ([]Task, error) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	if c.err != nil {
+		return nil, c.err
+	}
+	return c.tasks(subsystem, recursive)
+}
+
+func (c *cgroup) tasks(subsystem Name, recursive bool) ([]Task, error) {
+	s := c.getSubsystem(subsystem)
+	sp, err := c.path(subsystem)
+	if err != nil {
+		return nil, err
+	}
+	path := s.(pather).Path(sp)
+	var tasks []Task
+	err = filepath.Walk(path, func(p string, info os.FileInfo, err error) error {
+		if err != nil {
+			return err
+		}
+		if !recursive && info.IsDir() {
+			if p == path {
+				return nil
+			}
+			return filepath.SkipDir
+		}
+		dir, name := filepath.Split(p)
+		if name != cgroupTasks {
+			return nil
+		}
+		procs, err := readTasksPids(dir, subsystem)
+		if err != nil {
+			return err
+		}
+		tasks = append(tasks, procs...)
+		return nil
+	})
+	return tasks, err
+}
+
 // Freeze freezes the entire cgroup and all the processes inside it
 func (c *cgroup) Freeze() error {
 	c.mu.Lock()
diff --git a/cgroup_test.go b/cgroup_test.go
@@ -159,6 +159,41 @@ func TestListPids(t *testing.T) {
 	}
 }
 
+func TestListTasksPids(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
+	}
+	if err := control.AddTask(Process{Pid: 1234}); err != nil {
+		t.Error(err)
+		return
+	}
+	for _, s := range Subsystems() {
+		if err := checkTaskid(mock, filepath.Join(string(s), "test"), 1234); err != nil {
+			t.Error(err)
+			return
+		}
+	}
+	tasks, err := control.Tasks(Freezer, false)
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if l := len(tasks); l != 1 {
+		t.Errorf("should have one task but received %d", l)
+		return
+	}
+	if tasks[0].Pid != 1234 {
+		t.Errorf("expected task pid %d but received %d", 1234, tasks[0].Pid)
+	}
+}
+
 func readValue(mock *mockCgroup, path string) (string, error) {
 	data, err := ioutil.ReadFile(filepath.Join(mock.root, path))
 	if err != nil {
diff --git a/control.go b/control.go
@@ -44,6 +44,15 @@ type Process struct {
 	Path string
 }
 
+type Task struct {
+	// Subsystem is the name of the subsystem that the task is in
+	Subsystem Name
+	// Pid is the process id of the task
+	Pid int
+	// Path is the full path of the subsystem and location that the task is in
+	Path string
+}
+
 // Cgroup handles interactions with the individual groups to perform
 // actions on them as them main interface to this cgroup package
 type Cgroup interface {
@@ -64,6 +73,8 @@ type Cgroup interface {
 	Update(resources *specs.LinuxResources) error
 	// Processes returns all the processes in a select subsystem for the cgroup
 	Processes(Name, bool) ([]Process, error)
+	// Tasks returns all the tasks in a select subsystem for the cgroup
+	Tasks(Name, bool) ([]Task, error)
 	// Freeze freezes or pauses all processes inside the cgroup
 	Freeze() error
 	// Thaw thaw or resumes all processes inside the cgroup
diff --git a/utils.go b/utils.go
@@ -111,7 +111,7 @@ func remove(path string) error {
 	return fmt.Errorf("cgroups: unable to remove path %q", path)
 }
 
-// readPids will read all the pids in a cgroup by the provided path
+// readPids will read all the pids of processes in a cgroup by the provided path
 func readPids(path string, subsystem Name) ([]Process, error) {
 	f, err := os.Open(filepath.Join(path, cgroupProcs))
 	if err != nil {
@@ -138,6 +138,33 @@ func readPids(path string, subsystem Name) ([]Process, error) {
 	return out, nil
 }
 
+// readTasksPids will read all the pids of tasks in a cgroup by the provided path
+func readTasksPids(path string, subsystem Name) ([]Task, error) {
+	f, err := os.Open(filepath.Join(path, cgroupTasks))
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+	var (
+		out []Task
+		s   = bufio.NewScanner(f)
+	)
+	for s.Scan() {
+		if t := s.Text(); t != "" {
+			pid, err := strconv.Atoi(t)
+			if err != nil {
+				return nil, err
+			}
+			out = append(out, Task{
+				Pid:       pid,
+				Subsystem: subsystem,
+				Path:      path,
+			})
+		}
+	}
+	return out, nil
+}
+
 func hugePageSizes() ([]string, error) {
 	var (
 		pageSizes []string
