From 722de8ddcc875c7e8e9a228f9d88454bae31fd40 Mon Sep 17 00:00:00 2001
From: Alex Wied <centromere@users.noreply.github.com>
Date: Tue, 19 Jul 2022 02:09:46 -0400
Subject: [PATCH] libstore/globals.cc: Move cgroup detection to libutil

---
 src/libstore/globals.cc | 54 +++++------------------------------------
 src/libutil/util.cc     | 51 ++++++++++++++++++++++++++++++++++++++
 src/libutil/util.hh     |  3 +++
 3 files changed, 60 insertions(+), 48 deletions(-)

diff --git a/src/libstore/globals.cc b/src/libstore/globals.cc
index 48df839fa..d724897bb 100644
--- a/src/libstore/globals.cc
+++ b/src/libstore/globals.cc
@@ -11,11 +11,6 @@
 #include <dlfcn.h>
 #include <sys/utsname.h>
 
-#if __linux__
-#include <mntent.h>
-#include <cmath>
-#endif
-
 #include <nlohmann/json.hpp>
 
 
@@ -119,50 +114,13 @@ std::vector<Path> getUserConfigFiles()
 
 unsigned int Settings::getDefaultCores()
 {
-    unsigned int concurrency = std::max(1U, std::thread::hardware_concurrency());
+    const unsigned int concurrency = std::max(1U, std::thread::hardware_concurrency());
+    const unsigned int maxCPU = getMaxCPU();
 
-    #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;
+    if (maxCPU > 0)
+      return maxCPU;
+    else
+      return concurrency;
 }
 
 StringSet Settings::getDefaultSystemFeatures()
diff --git a/src/libutil/util.cc b/src/libutil/util.cc
index 28df30fef..be6fe091f 100644
--- 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
index d3ed15b0b..29227ecc6 100644
--- 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. */