Move daemon process into sub-cgroup

src/libstore/unix/build/local-derivation-goal.cc

@@ -444,25 +444,22 @@ void LocalDerivationGoal::startBuilder()
         #if __linux__
         experimentalFeatureSettings.require(Xp::Cgroups);
 
+        /* If we're running from the daemon, then this will return the
+           root cgroup of the service. Otherwise, it will return the
+           current cgroup. */
+        auto rootCgroup = getRootCgroup();
         auto cgroupFS = getCgroupFS();
         if (!cgroupFS)
             throw Error("cannot determine the cgroups file system");
-
-        auto ourCgroups = getCgroups("/proc/self/cgroup");
-        auto ourCgroup = ourCgroups[""];
-        if (ourCgroup == "")
-            throw Error("cannot determine cgroup name from /proc/self/cgroup");
-
-        auto ourCgroupPath = canonPath(*cgroupFS + "/" + ourCgroup);
-
-        if (!pathExists(ourCgroupPath))
-            throw Error("expected cgroup directory '%s'", ourCgroupPath);
+        auto rootCgroupPath = canonPath(*cgroupFS + "/" + rootCgroup);
+        if (!pathExists(rootCgroupPath))
+            throw Error("expected cgroup directory '%s'", rootCgroupPath);
 
         static std::atomic<unsigned int> counter{0};
 
         cgroup = buildUser
-            ? fmt("%s/nix-build-uid-%d", ourCgroupPath, buildUser->getUID())
-            : fmt("%s/nix-build-pid-%d-%d", ourCgroupPath, getpid(), counter++);
+            ? fmt("%s/nix-build-uid-%d", rootCgroupPath, buildUser->getUID())
+            : fmt("%s/nix-build-pid-%d-%d", rootCgroupPath, getpid(), counter++);
 
         debug("using cgroup '%s'", *cgroup);
 

src/libutil/current-process.cc

@@ -32,11 +32,7 @@ unsigned int getMaxCPU()
         auto cgroupFS = getCgroupFS();
         if (!cgroupFS) return 0;
 
-        auto cgroups = getCgroups("/proc/self/cgroup");
-        auto cgroup = cgroups[""];
-        if (cgroup == "") return 0;
-
-        auto cpuFile = *cgroupFS + "/" + cgroup + "/cpu.max";
+        auto cpuFile = *cgroupFS + "/" + getCurrentCgroup() + "/cpu.max";
 
         auto cpuMax = readFile(cpuFile);
         auto cpuMaxParts = tokenizeString<std::vector<std::string>>(cpuMax, " \n");

src/libutil/linux/cgroup.cc

@@ -144,4 +144,23 @@ CgroupStats destroyCgroup(const Path & cgroup)
     return destroyCgroup(cgroup, true);
 }
 
+std::string getCurrentCgroup()
+{
+    auto cgroupFS = getCgroupFS();
+    if (!cgroupFS)
+        throw Error("cannot determine the cgroups file system");
+
+    auto ourCgroups = getCgroups("/proc/self/cgroup");
+    auto ourCgroup = ourCgroups[""];
+    if (ourCgroup == "")
+        throw Error("cannot determine cgroup name from /proc/self/cgroup");
+    return ourCgroup;
+}
+
+std::string getRootCgroup()
+{
+    static std::string rootCgroup = getCurrentCgroup();
+    return rootCgroup;
+}
+
 }

src/libutil/linux/cgroup.hh

@@ -25,4 +25,13 @@ struct CgroupStats
  */
 CgroupStats destroyCgroup(const Path & cgroup);
 
+std::string getCurrentCgroup();
+
+/**
+ * Get the cgroup that should be used as the parent when creating new
+ * sub-cgroups. The first time this is called, the current cgroup will be
+ * returned, and then all subsequent calls will return the original cgroup.
+ */
+std::string getRootCgroup();
+
 }

src/nix/unix/daemon.cc

@@ -33,6 +33,10 @@
 #include <grp.h>
 #include <fcntl.h>
 
+#if __linux__
+#include "cgroup.hh"
+#endif
+
 #if __APPLE__ || __FreeBSD__
 #include <sys/ucred.h>
 #endif
@@ -312,6 +316,27 @@ static void daemonLoop(std::optional<TrustedFlag> forceTrustClientOpt)
     //  Get rid of children automatically; don't let them become zombies.
     setSigChldAction(true);
 
+    #if __linux__
+    if (settings.useCgroups) {
+        experimentalFeatureSettings.require(Xp::Cgroups);
+
+        //  This also sets the root cgroup to the current one.
+        auto rootCgroup = getRootCgroup();
+        auto cgroupFS = getCgroupFS();
+        if (!cgroupFS)
+            throw Error("cannot determine the cgroups file system");
+        auto rootCgroupPath = canonPath(*cgroupFS + "/" + rootCgroup);
+        if (!pathExists(rootCgroupPath))
+            throw Error("expected cgroup directory '%s'", rootCgroupPath);
+        auto daemonCgroupPath = rootCgroupPath + "/nix-daemon";
+        //  Create new sub-cgroup for the daemon.
+        if (mkdir(daemonCgroupPath.c_str(), 0755) != 0 && errno != EEXIST)
+            throw SysError("creating cgroup '%s'", daemonCgroupPath);
+        //  Move daemon into the new cgroup.
+        writeFile(daemonCgroupPath + "/cgroup.procs", fmt("%d", getpid()));
+    }
+    #endif
+
     //  Loop accepting connections.
     while (1) {
 

tests/nixos/cgroups/default.nix

@@ -0,0 +1,40 @@
+{ nixpkgs, ... }:
+
+{
+  name = "cgroups";
+
+  nodes =
+    {
+      host =
+        { config, pkgs, ... }:
+        { virtualisation.additionalPaths = [ pkgs.stdenvNoCC ];
+          nix.extraOptions =
+            ''
+              extra-experimental-features = nix-command auto-allocate-uids cgroups
+              extra-system-features = uid-range
+            '';
+          nix.settings.use-cgroups = true;
+          nix.nixPath = [ "nixpkgs=${nixpkgs}" ];
+        };
+    };
+
+  testScript = { nodes }: ''
+    start_all()
+
+    host.wait_for_unit("multi-user.target")
+
+    # Start build in background
+    host.execute("NIX_REMOTE=daemon nix build --auto-allocate-uids --file ${./hang.nix} >&2 &")
+    service = "/sys/fs/cgroup/system.slice/nix-daemon.service"
+
+    # Wait for cgroups to be created
+    host.succeed(f"until [ -e {service}/nix-daemon ]; do sleep 1; done", timeout=30)
+    host.succeed(f"until [ -e {service}/nix-build-uid-* ]; do sleep 1; done", timeout=30)
+
+    # Check that there aren't processes where there shouldn't be, and that there are where there should be
+    host.succeed(f'[ -z "$(cat {service}/cgroup.procs)" ]')
+    host.succeed(f'[ -n "$(cat {service}/nix-daemon/cgroup.procs)" ]')
+    host.succeed(f'[ -n "$(cat {service}/nix-build-uid-*/cgroup.procs)" ]')
+  '';
+
+}

tests/nixos/cgroups/hang.nix

@@ -0,0 +1,10 @@
+{ }:
+
+with import <nixpkgs> {};
+
+runCommand "hang"
+  { requiredSystemFeatures = "uid-range";
+  }
+  ''
+    sleep infinity
+  ''

tests/nixos/default.nix

@@ -157,4 +157,6 @@ in
   s3-binary-cache-store = runNixOSTestFor "x86_64-linux" ./s3-binary-cache-store.nix;
 
   fsync = runNixOSTestFor "x86_64-linux" ./fsync.nix;
+
+  cgroups = runNixOSTestFor "x86_64-linux" ./cgroups;
 }