Merge pull request #2591 from SR4ven/cgroup2

Add support for cgroup2 in CPU quota parsing

src/util.cc

@@ -39,6 +39,7 @@
 #include <sys/time.h>
 #endif
 
+#include <algorithm>
 #include <vector>
 
 #if defined(__APPLE__) || defined(__FreeBSD__)
@@ -689,16 +690,33 @@ map<string, string> ParseMountInfo(map<string, CGroupSubSys>& subsystems) {
     MountPoint mp;
     if (!mp.parse(line))
       continue;
-    if (mp.fsType != "cgroup")
-      continue;
-    for (size_t i = 0; i < mp.superOptions.size(); i++) {
-      string opt = mp.superOptions[i].AsString();
-      map<string, CGroupSubSys>::iterator subsys = subsystems.find(opt);
-      if (subsys == subsystems.end())
+    if (mp.fsType == "cgroup") {
+      for (size_t i = 0; i < mp.superOptions.size(); i++) {
+        std::string opt = mp.superOptions[i].AsString();
+        auto subsys = subsystems.find(opt);
+        if (subsys == subsystems.end()) {
+          continue;
+        }
+        std::string newPath = mp.translate(subsys->second.name);
+        if (!newPath.empty()) {
+          cgroups.emplace(opt, newPath);
+        }
+      }
+    } else if (mp.fsType == "cgroup2") {
+      // Find cgroup2 entry in format "0::/path/to/cgroup"
+      auto subsys = std::find_if(subsystems.begin(), subsystems.end(),
+                                 [](const auto& sys) {
+                                   return sys.first == "" && sys.second.id == 0;
+                                 });
+      if (subsys == subsystems.end()) {
         continue;
-      string newPath = mp.translate(subsys->second.name);
-      if (!newPath.empty())
-        cgroups.insert(make_pair(opt, newPath));
+      }
+      std::string path = mp.mountPoint.AsString();
+      if (subsys->second.name != "/") {
+        // Append the relative path for the cgroup to the mount point
+        path.append(subsys->second.name);
+      }
+      cgroups.emplace("cgroup2", path);
     }
   }
   return cgroups;
@@ -722,23 +740,73 @@ map<string, CGroupSubSys> ParseSelfCGroup() {
   return cgroups;
 }
 
-int ParseCPUFromCGroup() {
-  map<string, CGroupSubSys> subsystems = ParseSelfCGroup();
-  map<string, string> cgroups = ParseMountInfo(subsystems);
-  map<string, string>::iterator cpu = cgroups.find("cpu");
-  if (cpu == cgroups.end())
-    return -1;
-  std::pair<int64_t, bool> quota = readCount(cpu->second + "/cpu.cfs_quota_us");
+int ParseCgroupV1(std::string& path) {
+  std::pair<int64_t, bool> quota = readCount(path + "/cpu.cfs_quota_us");
   if (!quota.second || quota.first == -1)
     return -1;
-  std::pair<int64_t, bool> period =
-      readCount(cpu->second + "/cpu.cfs_period_us");
+  std::pair<int64_t, bool> period = readCount(path + "/cpu.cfs_period_us");
   if (!period.second)
     return -1;
   if (period.first == 0)
     return -1;
   return quota.first / period.first;
 }
+
+int ParseCgroupV2(std::string& path) {
+  // Read CPU quota from cgroup v2
+  std::ifstream cpu_max(path + "/cpu.max");
+  if (!cpu_max.is_open()) {
+    return -1;
+  }
+  std::string max_line;
+  if (!std::getline(cpu_max, max_line) || max_line.empty()) {
+    return -1;
+  }
+  // Format is "quota period" or "max period"
+  size_t space_pos = max_line.find(' ');
+  if (space_pos == string::npos) {
+    return -1;
+  }
+  std::string quota_str = max_line.substr(0, space_pos);
+  std::string period_str = max_line.substr(space_pos + 1);
+  if (quota_str == "max") {
+    return -1;  // No CPU limit set
+  }
+  // Convert quota string to integer
+  char* quota_end = nullptr;
+  errno = 0;
+  int64_t quota = strtoll(quota_str.c_str(), &quota_end, 10);
+  // Check for conversion errors
+  if (errno == ERANGE || quota_end == quota_str.c_str() || *quota_end != '\0' ||
+      quota <= 0) {
+    return -1;
+  }
+  // Convert period string to integer
+  char* period_end = nullptr;
+  errno = 0;
+  int64_t period = strtoll(period_str.c_str(), &period_end, 10);
+  // Check for conversion errors
+  if (errno == ERANGE || period_end == period_str.c_str() ||
+      *period_end != '\0' || period <= 0) {
+    return -1;
+  }
+  return quota / period;
+}
+
+int ParseCPUFromCGroup() {
+  auto subsystems = ParseSelfCGroup();
+  auto cgroups = ParseMountInfo(subsystems);
+
+  // Prefer cgroup v2 if both v1 and v2 should be present
+  if (const auto cgroup2 = cgroups.find("cgroup2"); cgroup2 != cgroups.end()) {
+    return ParseCgroupV2(cgroup2->second);
+  }
+
+  if (const auto cpu = cgroups.find("cpu"); cpu != cgroups.end()) {
+    return ParseCgroupV1(cpu->second);
+  }
+  return -1;
+}
 #endif
 
 int GetProcessorCount() {