feat: support cpu limit by cgroups on linux

1 files changed, 161 insertions, 2 deletions

diff --git a/src/util.cc b/src/util.cc
index 2e85513..287debc 100644
--- a/src/util.cc
+++ b/src/util.cc
@@ -49,6 +49,9 @@
 #include <libperfstat.h>
 #elif defined(linux) || defined(__GLIBC__)
 #include <sys/sysinfo.h>
+#include <fstream>
+#include <map>
+#include "string_piece_util.h"
 #endif
 
 #if defined(__FreeBSD__)
@@ -498,6 +501,155 @@ string StripAnsiEscapeCodes(const string& in) {
   return stripped;
 }
 
+#if defined(linux) || defined(__GLIBC__)
+std::pair<int64_t, bool> readCount(const std::string& path) {
+  std::ifstream file(path.c_str());
+  if (!file.is_open())
+    return std::make_pair(0, false);
+  int64_t n = 0;
+  file >> n;
+  if (file.good())
+    return std::make_pair(n, true);
+  return std::make_pair(0, false);
+}
+
+struct MountPoint {
+  int mountId;
+  int parentId;
+  StringPiece deviceId;
+  StringPiece root;
+  StringPiece mountPoint;
+  vector<StringPiece> options;
+  vector<StringPiece> optionalFields;
+  StringPiece fsType;
+  StringPiece mountSource;
+  vector<StringPiece> superOptions;
+  bool parse(const string& line) {
+    vector<StringPiece> pieces = SplitStringPiece(line, ' ');
+    if (pieces.size() < 10)
+      return false;
+    size_t optionalStart = 0;
+    for (size_t i = 6; i < pieces.size(); i++) {
+      if (pieces[i] == "-") {
+        optionalStart = i + 1;
+        break;
+      }
+    }
+    if (optionalStart == 0)
+      return false;
+    if (optionalStart + 3 != pieces.size())
+      return false;
+    mountId = atoi(pieces[0].AsString().c_str());
+    parentId = atoi(pieces[1].AsString().c_str());
+    deviceId = pieces[2];
+    root = pieces[3];
+    mountPoint = pieces[4];
+    options = SplitStringPiece(pieces[5], ',');
+    optionalFields =
+        vector<StringPiece>(&pieces[6], &pieces[optionalStart - 1]);
+    fsType = pieces[optionalStart];
+    mountSource = pieces[optionalStart + 1];
+    superOptions = SplitStringPiece(pieces[optionalStart + 2], ',');
+    return true;
+  }
+  string translate(string& path) const {
+    // path must be sub dir of root
+    if (path.compare(0, root.len_, root.str_, root.len_) != 0) {
+      return string();
+    }
+    path.erase(0, root.len_);
+    if (path == ".." || (path.length() > 2 && path.compare(0, 3, "../") == 0)) {
+      return string();
+    }
+    return mountPoint.AsString() + "/" + path;
+  }
+};
+
+struct CGroupSubSys {
+  int id;
+  string name;
+  vector<string> subsystems;
+  bool parse(string& line) {
+    size_t first = line.find(':');
+    if (first == string::npos)
+      return false;
+    line[first] = '\0';
+    size_t second = line.find(':', first + 1);
+    if (second == string::npos)
+      return false;
+    line[second] = '\0';
+    id = atoi(line.c_str());
+    name = line.substr(second + 1);
+    vector<StringPiece> pieces =
+        SplitStringPiece(StringPiece(line.c_str() + first + 1), ',');
+    for (size_t i = 0; i < pieces.size(); i++) {
+      subsystems.push_back(pieces[i].AsString());
+    }
+    return true;
+  }
+};
+
+map<string, string> ParseMountInfo(map<string, CGroupSubSys>& subsystems) {
+  map<string, string> cgroups;
+  ifstream mountinfo("/proc/self/mountinfo");
+  if (!mountinfo.is_open())
+    return cgroups;
+  while (!mountinfo.eof()) {
+    string line;
+    getline(mountinfo, line);
+    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())
+        continue;
+      string newPath = mp.translate(subsys->second.name);
+      if (!newPath.empty())
+        cgroups.insert(make_pair(opt, newPath));
+    }
+  }
+  return cgroups;
+}
+
+map<string, CGroupSubSys> ParseSelfCGroup() {
+  map<string, CGroupSubSys> cgroups;
+  ifstream cgroup("/proc/self/cgroup");
+  if (!cgroup.is_open())
+    return cgroups;
+  string line;
+  while (!cgroup.eof()) {
+    getline(cgroup, line);
+    CGroupSubSys subsys;
+    if (!subsys.parse(line))
+      continue;
+    for (size_t i = 0; i < subsys.subsystems.size(); i++) {
+      cgroups.insert(make_pair(subsys.subsystems[i], subsys));
+    }
+  }
+  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");
+  if (!quota.second || quota.first == -1)
+    return -1;
+  std::pair<int64_t, bool> period =
+      readCount(cpu->second + "/cpu.cfs_period_us");
+  if (!period.second)
+    return -1;
+  return quota.first / period.first;
+}
+#endif
+
 int GetProcessorCount() {
 #ifdef _WIN32
   DWORD cpuCount = 0;
@@ -545,6 +697,11 @@ int GetProcessorCount() {
   }
   return cpuCount;
 #else
+  int cgroupCount = -1;
+  int schedCount = -1;
+#if defined(linux) || defined(__GLIBC__)
+  cgroupCount = ParseCPUFromCGroup();
+#endif
   // The number of exposed processors might not represent the actual number of
   // processors threads can run on. This happens when a CPU set limitation is
   // active, see https://github.com/ninja-build/ninja/issues/1278
@@ -558,10 +715,12 @@ int GetProcessorCount() {
 #elif defined(CPU_COUNT)
   cpu_set_t set;
   if (sched_getaffinity(getpid(), sizeof(set), &set) == 0) {
-    return CPU_COUNT(&set);
+    schedCount = CPU_COUNT(&set);
   }
 #endif
-  return sysconf(_SC_NPROCESSORS_ONLN);
+  if (cgroupCount >= 0 && schedCount >= 0) return std::min(cgroupCount, schedCount);
+  if (cgroupCount < 0 && schedCount < 0) return sysconf(_SC_NPROCESSORS_ONLN);
+  return std::max(cgroupCount, schedCount);
 #endif
 }