From 1af5d798a4d10a07e995d420a759f2fe752b583a Mon Sep 17 00:00:00 2001
From: Alex Wied <centromere@users.noreply.github.com>
Date: Fri, 15 Jul 2022 00:14:29 -0400
Subject: [PATCH] libstore/globals.cc: Automatically set cores based on cgroup
 CPU limit

By default, Nix sets the "cores" setting to the number of CPUs which are
physically present on the machine. If cgroups are used to limit the CPU
and memory consumption of a large Nix build, the OOM killer may be
invoked.

For example, consider a GitLab CI pipeline which builds a large software
package. The GitLab runner spawns a container whose CPU is limited to 4
cores and whose memory is limited to 16 GiB. If the underlying machine
has 64 cores, Nix will invoke the build with -j64. In many cases, that
level of parallelism will invoke the OOM killer and the build will
completely fail.

This change sets the default value of "cores" to be
ceil(cpu_quota / cpu_period), with a fallback to
std::thread::hardware_concurrency() if cgroups v2 is not detected.
---
 src/libstore/globals.cc | 50 ++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

diff --git a/src/libstore/globals.cc b/src/libstore/globals.cc
index 0f2ca4b15..48df839fa 100644
--- a/src/libstore/globals.cc
+++ b/src/libstore/globals.cc
@@ -11,6 +11,11 @@
 #include <dlfcn.h>
 #include <sys/utsname.h>
 
+#if __linux__
+#include <mntent.h>
+#include <cmath>
+#endif
+
 #include <nlohmann/json.hpp>
 
 
@@ -114,7 +119,50 @@ std::vector<Path> getUserConfigFiles()
 
 unsigned int Settings::getDefaultCores()
 {
-    return std::max(1U, std::thread::hardware_concurrency());
+    unsigned int concurrency = std::max(1U, std::thread::hardware_concurrency());
+
+    #if __linux__
+    FILE *fp = fopen("/proc/mounts", "r");
+    if (!fp)
+        return concurrency;
+
+    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")
+                concurrency = std::ceil(std::stoi(quota) / std::stof(period));
+        }
+    }
+    #endif
+
+    return concurrency;
 }
 
 StringSet Settings::getDefaultSystemFeatures()