Merge pull request #1827 from okhowang/cpu-limit

Support cpu limit in docker

2 files changed, 197 insertions, 5 deletions

diff --git a/src/ninja.cc b/src/ninja.cc
index aea430d..2b71eb1 100644
--- a/src/ninja.cc
+++ b/src/ninja.cc
@@ -1394,11 +1394,28 @@ int ExceptionFilter(unsigned int code, struct _EXCEPTION_POINTERS *ep) {
 
 #endif  // _MSC_VER
 
+class DeferGuessParallelism {
+ public:
+  bool needGuess;
+  BuildConfig* config;
+
+  DeferGuessParallelism(BuildConfig* config)
+      : config(config), needGuess(true) {}
+
+  void Refresh() {
+    if (needGuess) {
+      needGuess = false;
+      config->parallelism = GuessParallelism();
+    }
+  }
+  ~DeferGuessParallelism() { Refresh(); }
+};
+
 /// Parse argv for command-line options.
 /// Returns an exit code, or -1 if Ninja should continue.
 int ReadFlags(int* argc, char*** argv,
               Options* options, BuildConfig* config) {
-  config->parallelism = GuessParallelism();
+  DeferGuessParallelism deferGuessParallelism(config);
 
   enum { OPT_VERSION = 1, OPT_QUIET = 2 };
   const option kLongOptions[] = {
@@ -1430,6 +1447,7 @@ int ReadFlags(int* argc, char*** argv,
         // We want to run N jobs in parallel. For N = 0, INT_MAX
         // is close enough to infinite for most sane builds.
         config->parallelism = value > 0 ? value : INT_MAX;
+        deferGuessParallelism.needGuess = false;
         break;
       }
       case 'k': {
@@ -1478,6 +1496,7 @@ int ReadFlags(int* argc, char*** argv,
         return 0;
       case 'h':
       default:
+        deferGuessParallelism.Refresh();
         Usage(*config);
         return 1;
     }
diff --git a/src/util.cc b/src/util.cc
index a2a0f27..483f4a6 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__)
@@ -499,8 +502,158 @@ 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;
 #ifndef _WIN64
   // Need to use GetLogicalProcessorInformationEx to get real core count on
   // machines with >64 cores. See https://stackoverflow.com/a/31209344/21475
@@ -525,13 +678,31 @@ int GetProcessorCount() {
         i += info->Size;
       }
       if (cores != 0) {
-        return cores;
+        cpuCount = cores;
       }
     }
   }
 #endif
-  return GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
+  if (cpuCount == 0) {
+    cpuCount = GetActiveProcessorCount(ALL_PROCESSOR_GROUPS);
+  }
+  JOBOBJECT_CPU_RATE_CONTROL_INFORMATION info;
+  // reference:
+  // https://docs.microsoft.com/en-us/windows/win32/api/winnt/ns-winnt-jobobject_cpu_rate_control_information
+  if (QueryInformationJobObject(NULL, JobObjectCpuRateControlInformation, &info,
+                                sizeof(info), NULL)) {
+    if (info.ControlFlags & (JOB_OBJECT_CPU_RATE_CONTROL_ENABLE |
+                             JOB_OBJECT_CPU_RATE_CONTROL_HARD_CAP)) {
+      return cpuCount * info.CpuRate / 10000;
+    }
+  }
+  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
@@ -545,10 +716,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
 }