Merge pull request #312 from Rouzip/cgroup2_thread

AkihiroSuda · web-flow · commit 5e7266aa5dc8 · 2023-12-29T20:38:10.000+09:00
feat: add Threads for cgroupv2
diff --git a/cgroup2/manager.go b/cgroup2/manager.go
@@ -21,6 +21,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io/fs"
 	"math"
 	"os"
 	"path/filepath"
@@ -478,9 +479,9 @@ func (c *Manager) Delete() error {
 	return remove(c.path)
 }
 
-func (c *Manager) Procs(recursive bool) ([]uint64, error) {
-	var processes []uint64
-	err := filepath.Walk(c.path, func(p string, info os.FileInfo, err error) error {
+func (c *Manager) getTasks(recursive bool, tType string) ([]uint64, error) {
+	var tasks []uint64
+	err := filepath.Walk(c.path, func(p string, info fs.FileInfo, err error) error {
 		if err != nil {
 			return err
 		}
@@ -491,17 +492,25 @@ func (c *Manager) Procs(recursive bool) ([]uint64, error) {
 			return filepath.SkipDir
 		}
 		_, name := filepath.Split(p)
-		if name != cgroupProcs {
+		if name != tType {
 			return nil
 		}
-		procs, err := parseCgroupProcsFile(p)
+		curTasks, err := parseCgroupTasksFile(p)
 		if err != nil {
 			return err
 		}
-		processes = append(processes, procs...)
+		tasks = append(tasks, curTasks...)
 		return nil
 	})
-	return processes, err
+	return tasks, err
+}
+
+func (c *Manager) Procs(recursive bool) ([]uint64, error) {
+	return c.getTasks(recursive, cgroupProcs)
+}
+
+func (c *Manager) Threads(recursive bool) ([]uint64, error) {
+	return c.getTasks(recursive, cgroupThreads)
 }
 
 func (c *Manager) MoveTo(destination *Manager) error {
diff --git a/cgroup2/manager_test.go b/cgroup2/manager_test.go
@@ -211,6 +211,9 @@ func TestKill(t *testing.T) {
 	pids, err := manager.Procs(true)
 	require.NoError(t, err)
 	require.Len(t, pids, numProcs, "pid count unexpected")
+	threads, err := manager.Threads(true)
+	require.NoError(t, err)
+	require.Len(t, threads, numProcs, "pid count unexpected")
 
 	// Now run kill, and check that nothing is running after.
 	err = manager.Kill()
diff --git a/cgroup2/utils.go b/cgroup2/utils.go
@@ -68,8 +68,9 @@ func remove(path string) error {
 	return fmt.Errorf("cgroups: unable to remove path %q: %w", path, err)
 }
 
-// parseCgroupProcsFile parses /sys/fs/cgroup/$GROUPPATH/cgroup.procs
-func parseCgroupProcsFile(path string) ([]uint64, error) {
+// parseCgroupTasksFile parses /sys/fs/cgroup/$GROUPPATH/cgroup.procs or
+// /sys/fs/cgroup/$GROUPPATH/cgroup.threads
+func parseCgroupTasksFile(path string) ([]uint64, error) {
 	f, err := os.Open(path)
 	if err != nil {
 		return nil, err
