Merge pull request NixOS#6813 from centromere/cgroup-cpu-detection

libstore/globals.cc: Automatically set cores based on cgroup CPU limit
rti · Jul 22, 2022 · e10807c · e10807c
2 parents fbd0a6c + 722de8d
commit e10807c
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 1 deletion.
diff --git a/src/libstore/globals.cc b/src/libstore/globals.cc
@@ -114,7 +114,13 @@ std::vector<Path> getUserConfigFiles()
 
 unsigned int Settings::getDefaultCores()
 {
-    return std::max(1U, std::thread::hardware_concurrency());
+    const unsigned int concurrency = std::max(1U, std::thread::hardware_concurrency());
+    const unsigned int maxCPU = getMaxCPU();
+
+    if (maxCPU > 0)
+      return maxCPU;
+    else
+      return concurrency;
 }
 
 StringSet Settings::getDefaultSystemFeatures()

diff --git a/src/libutil/util.cc b/src/libutil/util.cc
@@ -35,6 +35,9 @@
 #ifdef __linux__
 #include <sys/prctl.h>
 #include <sys/resource.h>
+
+#include <mntent.h>
+#include <cmath>
 #endif
 
 
@@ -788,7 +791,55 @@ void drainFD(int fd, Sink & sink, bool block)
     }
 }
 
+//////////////////////////////////////////////////////////////////////
+
+unsigned int getMaxCPU()
+{
+    #if __linux__
+    try {
+        FILE *fp = fopen("/proc/mounts", "r");
+        if (!fp)
+            return 0;
 
+        Strings cgPathParts;
+
+        struct mntent *ent;
+        while ((ent = getmntent(fp))) {
+            std::string mountType, mountPath;
+
+            mountType = ent->mnt_type;
+            mountPath = ent->mnt_dir;
+
+            if (mountType == "cgroup2") {
+                cgPathParts.push_back(mountPath);
+                break;
+            }
+        }
+
+        fclose(fp);
+
+        if (cgPathParts.size() > 0 && pathExists("/proc/self/cgroup")) {
+            std::string currentCgroup = readFile("/proc/self/cgroup");
+            Strings cgValues = tokenizeString<Strings>(currentCgroup, ":");
+            cgPathParts.push_back(trim(cgValues.back(), "\n"));
+            cgPathParts.push_back("cpu.max");
+            std::string fullCgPath = canonPath(concatStringsSep("/", cgPathParts));
+
+            if (pathExists(fullCgPath)) {
+                std::string cpuMax = readFile(fullCgPath);
+                std::vector<std::string> cpuMaxParts = tokenizeString<std::vector<std::string>>(cpuMax, " ");
+                std::string quota = cpuMaxParts[0];
+                std::string period = trim(cpuMaxParts[1], "\n");
+
+                if (quota != "max")
+                    return std::ceil(std::stoi(quota) / std::stof(period));
+            }
+        }
+    } catch (Error &) { ignoreException(); }
+    #endif
+
+    return 0;
+}
 
 //////////////////////////////////////////////////////////////////////
 

diff --git a/src/libutil/util.hh b/src/libutil/util.hh
@@ -182,6 +182,9 @@ std::string drainFD(int fd, bool block = true, const size_t reserveSize=0);
 
 void drainFD(int fd, Sink & sink, bool block = true);
 
+/* If cgroups are active, attempt to calculate the number of CPUs available.
+   If cgroups are unavailable or if cpu.max is set to "max", return 0. */
+unsigned int getMaxCPU();
 
 /* Automatic cleanup of resources. */