Ensure cgroup v2 cpu controller is enabled

Author: Furisto

components/ws-daemon/pkg/cpulimit/dispatch.go

@@ -12,8 +12,6 @@ import (
 	"sync"
 	"time"
 
-	"github.com/opencontainers/runc/libcontainer/cgroups/fs2"
-	"github.com/opencontainers/runc/libcontainer/configs"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/xerrors"
@@ -234,7 +232,7 @@ func newCFSController(basePath, cgroupPath string) (CFSController, error) {
 
 	if err == nil {
 		fullPath := filepath.Join(basePath, cgroupPath)
-		if err := ensureControllerEnabled(fullPath, "cpu"); err != nil {
+		if err := ensureControllerEnabled(basePath, cgroupPath, "cpu"); err != nil {
 			return nil, err
 		}
 
@@ -244,19 +242,59 @@ func newCFSController(basePath, cgroupPath string) (CFSController, error) {
 	return nil, err
 }
 
-func ensureControllerEnabled(targetPath, controller string) error {
-	controllerFile := filepath.Join(targetPath, "cgroup.controllers")
-	controllers, err := os.ReadFile(controllerFile)
+func ensureControllerEnabled(basePath, cgroupPath, controller string) error {
+	targetPath := filepath.Join(basePath, cgroupPath)
+	if enabled, err := isCpuControllerEnabled(targetPath); err != nil || enabled {
+		return err
+	}
+
+	err := writeCpuController(basePath)
 	if err != nil {
 		return err
 	}
 
+	levelPath := basePath
+	cgroupPath = strings.TrimPrefix(cgroupPath, "/")
+	levels := strings.Split(cgroupPath, string(os.PathSeparator))
+	for _, l := range levels[:len(levels)-1] {
+		levelPath = filepath.Join(levelPath, l)
+		err = writeCpuController(levelPath)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func isCpuControllerEnabled(path string) (bool, error) {
+	controllerFile := filepath.Join(path, "cgroup.controllers")
+	controllers, err := os.ReadFile(controllerFile)
+	if err != nil {
+		return false, err
+	}
+
 	for _, ctrl := range strings.Fields(string(controllers)) {
-		if ctrl == controller {
+		if ctrl == "cpu" {
 			// controller is already activated
-			return nil
+			return true, nil
 		}
 	}
 
-	return fs2.CreateCgroupPath(targetPath, &configs.Cgroup{})
+	return false, nil
+}
+
+func writeCpuController(path string) error {
+	f, err := os.OpenFile(filepath.Join(path, "cgroup.subtree_control"), os.O_WRONLY, 0)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	_, err = f.Write([]byte("+cpu"))
+	if err != nil {
+		return err
+	}
+
+	return nil
 }

components/ws-daemon/pkg/cpulimit/dispatch_test.go

@@ -0,0 +1,85 @@
+package cpulimit
+
+import (
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+var cgroupPath = []string{"kubepods", "burstable", "pods234sdf", "234as8df34"}
+
+func createHierarchy(t *testing.T, cpuEnabled bool) (string, string) {
+	testRoot := t.TempDir()
+	if err := os.WriteFile(filepath.Join(testRoot, "cgroup.controllers"), []byte(""), 0755); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := os.WriteFile(filepath.Join(testRoot, "cgroup.subtree_control"), []byte(""), 0755); err != nil {
+		t.Fatal(err)
+	}
+
+	testCgroup := ""
+	for i, level := range cgroupPath {
+		testCgroup = filepath.Join(testCgroup, level)
+		fullPath := filepath.Join(testRoot, testCgroup)
+		if err := os.Mkdir(fullPath, 0o755); err != nil {
+			t.Fatal(err)
+		}
+
+		ctrlFile, err := os.Create(filepath.Join(fullPath, "cgroup.controllers"))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer ctrlFile.Close()
+
+		if cpuEnabled {
+			if _, err := ctrlFile.WriteString("cpu"); err != nil {
+				t.Fatal(err)
+			}
+		}
+
+		subTreeFile, err := os.Create(filepath.Join(fullPath, "cgroup.subtree_control"))
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer subTreeFile.Close()
+
+		if cpuEnabled && i < len(cgroupPath)-1 {
+			if _, err := subTreeFile.WriteString("cpu"); err != nil {
+				t.Fatal(err)
+			}
+		}
+	}
+
+	return testRoot, testCgroup
+}
+
+func TestEnableController(t *testing.T) {
+	root, cgroup := createHierarchy(t, false)
+	if err := ensureControllerEnabled(root, cgroup, "cpu"); err != nil {
+		t.Fatal(err)
+	}
+
+	levelPath := root
+	for _, level := range cgroupPath {
+		verifyCpuControllerToggled(t, levelPath, true)
+		levelPath = filepath.Join(levelPath, level)
+	}
+
+	verifyCpuControllerToggled(t, levelPath, false)
+}
+
+func verifyCpuControllerToggled(t *testing.T, path string, enabled bool) {
+	t.Helper()
+
+	content, err := os.ReadFile(filepath.Join(path, "cgroup.subtree_control"))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if enabled && string(content) != "+cpu" {
+		t.Fatalf("%s should have enabled cpu controller", path)
+	} else if !enabled && string(content) == "+cpu" {
+		t.Fatalf("%s should not have enabled cpu controller", path)
+	}
+}