From da50f7d9ddf912a6e4ad8714e5aff01e85c09b21 Mon Sep 17 00:00:00 2001
From: zhongjiawei <zhongjiawei1@huawei.com>
Date: Wed, 26 Jul 2023 11:37:21 +0800
Subject: [PATCH] runc:support set cpuset.prefer_cpus using --cpuset-cpus

---
 libcontainer/cgroups/fs/cpuset.go | 118 +++++++++++++++++++++++++++++-
 1 file changed, 114 insertions(+), 4 deletions(-)

diff --git a/libcontainer/cgroups/fs/cpuset.go b/libcontainer/cgroups/fs/cpuset.go
index 939b72c..ef9164b 100644
--- a/libcontainer/cgroups/fs/cpuset.go
+++ b/libcontainer/cgroups/fs/cpuset.go
@@ -26,16 +26,126 @@ func (s *CpusetGroup) Apply(path string, r *configs.Resources, pid int) error {
 	return s.ApplyDir(path, r, pid)
 }
 
-func (s *CpusetGroup) Set(path string, r *configs.Resources) error {
-	var ret error
-	if r.CpusetCpus != "" {
-		if err := cgroups.WriteFile(path, "cpuset.cpus", r.CpusetCpus); err != nil {
+func parseCpus(cpus string) (map[int]bool, error) {
+	persedCpus := make(map[int]bool)
+	for _, cpu := range strings.Split(cpus, ",") {
+		invalidFormatError := fmt.Errorf("invalid cpuset format: %s", cpu)
+		if cpuInt, err := strconv.Atoi(cpu); err == nil {
+			if cpuInt < 0 {
+				return nil, invalidFormatError
+			}
+			persedCpus[cpuInt] = true
+			continue
+		}
+
+		ranges := strings.Split(cpu, "-")
+		if len(ranges) != 2 {
+			return nil, invalidFormatError
+		}
+		lower, err1 := strconv.Atoi(ranges[0])
+		upper, err2 := strconv.Atoi(ranges[1])
+		if err1 != nil || err2 != nil {
+			return nil, invalidFormatError
+		}
+		if lower < 0 || upper < 0 {
+			return nil, invalidFormatError
+		}
+		if lower > upper {
+			return nil, invalidFormatError
+		}
+		const max = 8192
+		if upper > max {
+			return nil, fmt.Errorf("value of out range, maximum is %d", max)
+		}
+		for i := lower; i <= upper; i++ {
+			persedCpus[i] = true
+		}
+	}
+
+	return persedCpus, nil
+}
+
+func checkInCpuset(cpuset, prefer map[int]bool) bool {
+	for k := range prefer {
+		if _, ok := cpuset[k]; !ok {
+			return false
+		}
+	}
+	return true
+}
+
+func parseCpuset(cpuset string) (string, string, error) {
+	var cpusetStr, preferCpusStr string
+	invalidFormatError := fmt.Errorf("invalid cpuset format: %s", cpuset)
+
+	splits := strings.Split(cpuset, "+")
+	if len(splits) == 1 {
+		cpusetStr = cpuset
+	} else if len(splits) == 2 {
+		if !strings.HasSuffix(splits[0], ",") {
+			return "", "", invalidFormatError
+		}
+		cpusetStr = strings.TrimSuffix(splits[0], ",")
+		preferCpusStr = splits[1]
+
+		cpusetMap, err := parseCpus(cpusetStr)
+		if err != nil {
+			return "", "", err
+		}
+		preferCpusMap, err := parseCpus(preferCpusStr)
+		if err != nil {
+			return "", "", err
+		}
+		if !checkInCpuset(cpusetMap, preferCpusMap) {
+			return "", "", fmt.Errorf("invalid preferred_cpus: %s not in cpuset: %s", preferCpusStr, cpusetStr)
+		}
+	} else {
+		return "", "", invalidFormatError
+	}
+
+	return cpusetStr, preferCpusStr, nil
+}
+
+func (s *CpusetGroup) setCpuset(path, cpuset string) error {
+	if _, err := os.Stat("/proc/sys/kernel/sched_util_low_pct"); err != nil {
+		var ret error
+		if err := writeFile(path, "cpuset.cpus", cpuset); err != nil {
 			ret = fmt.Errorf("failed to set cpuset.cpus, %v", err)
 			if _, err := os.Stat(path); err != nil {
 				ret = fmt.Errorf("%v, failed to stat %v, %v", ret, path, err)
 			}
 			return ret
 		}
+		return nil
+	}
+
+	cpusetStr, preferCpusStr, err := parseCpuset(cpuset)
+	if err != nil {
+		return err
+	}
+
+	if _, err := os.Stat(path); err != nil {
+		return fmt.Errorf("failed to stat %v, %v", path, err)
+	}
+	if err := writeFile(path, "cpuset.preferred_cpus", ""); err != nil {
+		return fmt.Errorf("failed to set cpuset.preferred_cpus to nil, %v", err)
+	}
+	if err := writeFile(path, "cpuset.cpus", cpusetStr); err != nil {
+		return fmt.Errorf("failed to set cpuset.cpus, %v", err)
+	}
+	if preferCpusStr != "" {
+		if err := writeFile(path, "cpuset.preferred_cpus", preferCpusStr); err != nil {
+			return fmt.Errorf("failed to set cpuset.preferred_cpus, %v", err)
+		}
+	}
+
+	return nil
+}
+
+func (s *CpusetGroup) Set(path string, r *configs.Resources) error {
+	var ret error
+	if r.CpusetCpus != "" {
+		return s.setCpuset(path, cgroup.Resources.CpusetCpus)
 	}
 	if r.CpusetMems != "" {
 		if err := cgroups.WriteFile(path, "cpuset.mems", r.CpusetMems); err != nil {
-- 
2.33.0