From 570c443f560e015cf02e4f96102eaaa0e6853562 Mon Sep 17 00:00:00 2001
From: Eelco Dolstra <edolstra@gmail.com>
Date: Sat, 16 May 2020 21:21:41 +0200
Subject: [PATCH] Simplify cgroup creation

---
 src/libstore/build.cc  | 53 +++++++++++++-----------------------------
 src/libstore/cgroup.cc | 17 ++++++++++++++
 src/libstore/cgroup.hh |  2 ++
 3 files changed, 35 insertions(+), 37 deletions(-)

diff --git a/src/libstore/build.cc b/src/libstore/build.cc
index 1f1468d97..97554e9cf 100644
--- a/src/libstore/build.cc
+++ b/src/libstore/build.cc
@@ -2375,50 +2375,30 @@ void DerivationGoal::startBuilder()
 
 #if __linux__
     if (useChroot) {
-        /* Create a cgroup. */
+        /* Create a systemd cgroup since that's the minimum required
+           by systemd-nspawn. */
         // FIXME: do we want to use the parent cgroup? We should
         // always use the same cgroup regardless of whether we're the
         // daemon or run from a user session via sudo.
-        std::string msg;
-        std::vector<Path> cgroups;
-        for (auto & line : tokenizeString<std::vector<std::string>>(readFile("/proc/self/cgroup"), "\n")) {
-            static std::regex regex("([0-9]+):([^:]*):(.*)");
-            std::smatch match;
-            if (!std::regex_match(line, match, regex))
-                throw Error("invalid line '%s' in '/proc/self/cgroup'", line);
+        auto ourCgroups = getCgroups("/proc/self/cgroup");
+        auto systemdCgroup = ourCgroups["systemd"];
+        if (systemdCgroup == "")
+            throw Error("'systemd' cgroup does not exist");
 
-            /* We only create a systemd cgroup, since that's enough
-               for running systemd-nspawn. */
-            std::string name;
-            if (match[2] == "name=systemd")
-                name = "systemd";
-            //else if (match[2] == "")
-            //    name = "unified";
-            else continue;
+        auto hostCgroup = canonPath("/sys/fs/cgroup/systemd/" + systemdCgroup);
 
-            std::string cgroup = match[3];
+        if (!pathExists(hostCgroup))
+            throw Error("expected cgroup directory '%s'", hostCgroup);
 
-            auto hostCgroup = canonPath("/sys/fs/cgroup/" + name + "/" + cgroup);
+        auto childCgroup = fmt("%s/nix-%d", hostCgroup, buildUser->getUID());
 
-            if (!pathExists(hostCgroup))
-                throw Error("expected cgroup directory '%s'", hostCgroup);
+        destroyCgroup(childCgroup);
 
-            auto childCgroup = fmt("%s/nix-%d", hostCgroup, buildUser->getUID());
+        if (mkdir(childCgroup.c_str(), 0755) == -1)
+            throw SysError("creating cgroup '%s'", childCgroup);
 
-            destroyCgroup(childCgroup);
-
-            if (mkdir(childCgroup.c_str(), 0755) == -1)
-                throw SysError("creating cgroup '%s'", childCgroup);
-
-            chownToBuilder(childCgroup);
-            chownToBuilder(childCgroup + "/cgroup.procs");
-            if (name == "unified") {
-                chownToBuilder(childCgroup + "/cgroup.threads");
-                chownToBuilder(childCgroup + "/cgroup.subtree_control");
-            }
-
-            cgroups.push_back(childCgroup);
-        }
+        chownToBuilder(childCgroup);
+        chownToBuilder(childCgroup + "/cgroup.procs");
 
         /* Set up private namespaces for the build:
 
@@ -2545,8 +2525,7 @@ void DerivationGoal::startBuilder()
             throw SysError("getting sandbox mount namespace");
 
         /* Move the child into its own cgroup. */
-        for (auto & childCgroup : cgroups)
-            writeFile(childCgroup + "/cgroup.procs", fmt("%d", (pid_t) pid));
+        writeFile(childCgroup + "/cgroup.procs", fmt("%d", (pid_t) pid));
 
         /* Signal the builder that we've updated its user namespace. */
         writeFull(userNamespaceSync.writeSide.get(), "1");
diff --git a/src/libstore/cgroup.cc b/src/libstore/cgroup.cc
index 8cd682e68..887facdca 100644
--- a/src/libstore/cgroup.cc
+++ b/src/libstore/cgroup.cc
@@ -9,6 +9,23 @@
 
 namespace nix {
 
+std::map<std::string, std::string> getCgroups(const Path & cgroupFile)
+{
+    std::map<std::string, std::string> cgroups;
+
+    for (auto & line : tokenizeString<std::vector<std::string>>(readFile(cgroupFile), "\n")) {
+        static std::regex regex("([0-9]+):([^:]*):(.*)");
+        std::smatch match;
+        if (!std::regex_match(line, match, regex))
+            throw Error("invalid line '%s' in '%s'", line, cgroupFile);
+
+        std::string name = hasPrefix(match[2], "name=") ? std::string(match[2], 5) : match[2];
+        cgroups.insert_or_assign(name, match[3]);
+    }
+
+    return cgroups;
+}
+
 void destroyCgroup(const Path & cgroup)
 {
     for (auto & entry : readDirectory(cgroup)) {
diff --git a/src/libstore/cgroup.hh b/src/libstore/cgroup.hh
index c7b09398e..dc6758957 100644
--- a/src/libstore/cgroup.hh
+++ b/src/libstore/cgroup.hh
@@ -6,6 +6,8 @@
 
 namespace nix {
 
+std::map<std::string, std::string> getCgroups(const Path & cgroupFile);
+
 void destroyCgroup(const Path & cgroup);
 
 }