Refactor CgroupsInfo to decouple the implementations between v1 and v2.

Also include an InvalidCgroupsInfo implementation that allows us to represent a cgroup that was attempted to be created but not successful (rather than using null to represent it).

The general way for clients to use this is to create the CgroupsInfo with the methods `CgroupsInfo.getBlazeSpawnsCgroup()`, `CgroupsInfo.getBlazeSpawnsCgroup().createIndividualSpawnCgroup()`. Then check the returned cgroup instance with `exists()` to ensure that it is created (and not invalid) before continuing with other actions e.g. `addProcess()`.

PiperOrigin-RevId: 605269206
Change-Id: I69fe32e3087e37e0705b840e10f07f37160402d0
diff --git a/.bazelci/postsubmit.yml b/.bazelci/postsubmit.yml
index e39cd94..7c0b70b 100644
--- a/.bazelci/postsubmit.yml
+++ b/.bazelci/postsubmit.yml
@@ -191,6 +191,7 @@
       # C++ coverage is not supported on macOS yet.
       - "-//src/test/shell/bazel:bazel_cc_code_coverage_test"
       # MacOS does not have cgroups so it can't support hardened sandbox
+      - "-//src/test/java/com/google/devtools/build/lib/sandbox:CgroupsInfoTest"
       - "-//src/test/shell/integration:bazel_hardened_sandboxed_worker_test"
       # https://github.com/bazelbuild/bazel/issues/16526
       - "-//src/test/shell/bazel:starlark_repository_test"
@@ -249,6 +250,7 @@
       # C++ coverage is not supported on macOS yet.
       - "-//src/test/shell/bazel:bazel_cc_code_coverage_test"
       # MacOS does not have cgroups so it can't support hardened sandbox
+      - "-//src/test/java/com/google/devtools/build/lib/sandbox:CgroupsInfoTest"
       - "-//src/test/shell/integration:bazel_hardened_sandboxed_worker_test"
       # https://github.com/bazelbuild/bazel/issues/16525
       - "-//src/test/java/com/google/devtools/build/lib/buildtool:KeepGoingTest"
diff --git a/.bazelci/presubmit.yml b/.bazelci/presubmit.yml
index 77552a5..e13c4f1 100644
--- a/.bazelci/presubmit.yml
+++ b/.bazelci/presubmit.yml
@@ -194,6 +194,7 @@
       # C++ coverage is not supported on macOS yet.
       - "-//src/test/shell/bazel:bazel_cc_code_coverage_test"
       # MacOS does not have cgroups so it can't support hardened sandbox
+      - "-//src/test/java/com/google/devtools/build/lib/sandbox:CgroupsInfoTest"
       - "-//src/test/shell/integration:bazel_hardened_sandboxed_worker_test"
       # https://github.com/bazelbuild/bazel/issues/16526
       - "-//src/test/shell/bazel:starlark_repository_test"
@@ -253,6 +254,7 @@
       # C++ coverage is not supported on macOS yet.
       - "-//src/test/shell/bazel:bazel_cc_code_coverage_test"
       # MacOS does not have cgroups so it can't support hardened sandbox
+      - "-//src/test/java/com/google/devtools/build/lib/sandbox:CgroupsInfoTest"
       - "-//src/test/shell/integration:bazel_hardened_sandboxed_worker_test"
       # https://github.com/bazelbuild/bazel/issues/16525
       - "-//src/test/java/com/google/devtools/build/lib/buildtool:KeepGoingTest"
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/BUILD b/src/main/java/com/google/devtools/build/lib/sandbox/BUILD
index f0fda18..88714c3 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/BUILD
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/BUILD
@@ -195,6 +195,8 @@
     name = "linux_sandbox",
     srcs = [
         "CgroupsInfo.java",
+        "CgroupsInfoV1.java",
+        "CgroupsInfoV2.java",
         "LinuxSandboxedSpawnRunner.java",
         "LinuxSandboxedStrategy.java",
     ],
@@ -222,7 +224,6 @@
         "//src/main/java/com/google/devtools/build/lib/profiler",
         "//src/main/java/com/google/devtools/build/lib/shell",
         "//src/main/java/com/google/devtools/build/lib/util:os",
-        "//src/main/java/com/google/devtools/build/lib/util:pair",
         "//src/main/java/com/google/devtools/build/lib/vfs",
         "//src/main/java/com/google/devtools/build/lib/vfs:pathfragment",
         "//third_party:flogger",
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfo.java b/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfo.java
index 101b184..7f15d69 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfo.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfo.java
@@ -17,20 +17,21 @@
 import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.annotations.VisibleForTesting;
-import com.google.common.base.Joiner;
-import com.google.common.base.Splitter;
+import com.google.common.base.Suppliers;
 import com.google.common.flogger.GoogleLogger;
 import com.google.common.io.Files;
-import com.google.devtools.build.lib.util.Pair;
+import com.google.devtools.build.lib.util.OS;
 import java.io.File;
 import java.io.IOException;
 import java.util.List;
+import java.util.function.Supplier;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import javax.annotation.Nullable;
 
 /** This class manages cgroups directories for memory-limiting sandboxed processes. */
-public class CgroupsInfo {
+public abstract class CgroupsInfo {
+
   private static final GoogleLogger logger = GoogleLogger.forEnclosingClass();
 
   /**
@@ -42,98 +43,152 @@
   private static final Pattern CGROUPS_MOUNT_PATTERN =
       Pattern.compile("^cgroup(|2)\\s+(\\S*)\\s+cgroup2?\\s+(\\S*).*");
 
-  public static final String PROC_SELF_MOUNTS_PATH = "/proc/self/mounts";
-  public static final String PROC_SELF_CGROUP_PATH = "/proc/self/cgroup";
-  /** If non-null, this is a cgroups directory that sandboxes can put their directories into. */
-  @Nullable private static volatile CgroupsInfo instance;
+  private static final String PROC_SELF_MOUNTS_PATH = "/proc/self/mounts";
+  private static final String PROC_SELF_CGROUP_PATH = "/proc/self/cgroup";
 
-  private final boolean isCgroupsV2;
-  // This is the directory where the cgroup is in, any related files pertaining to limits / resource
-  // usage or child cgroups (nested directories) are found here.
-  private final File cgroupDir;
-  private final File mountPoint;
+  private static final CgroupsInfo rootCgroup = getRootCgroup(new File(PROC_SELF_MOUNTS_PATH));
 
-  private CgroupsInfo(boolean isCgroupsV2, File cgroupDir, File mountPoint) {
-    this.isCgroupsV2 = isCgroupsV2;
-    this.cgroupDir = cgroupDir;
-    this.mountPoint = mountPoint;
+  private static final Supplier<CgroupsInfo> blazeSpawnsCgroupSupplier =
+      Suppliers.memoize(CgroupsInfo::createBlazeSpawnsCgroup);
+
+  /** Returns whether the local machine supports cgroups. */
+  public static boolean isSupported() {
+    return OS.getCurrent() == OS.LINUX && getBlazeSpawnsCgroup().canWrite();
   }
 
   /**
-   * Creates a cgroups directory for Blaze to place sandboxes in. Figures out whether cgroups v1 or
-   * v2 is available, and for cgroups v2 sets subtree control for the <code>memory</code> and <code>
-   * pids</code> controllers.
+   * Returns an instance of the root cgroup of the hierarchy, {@link InvalidCgroupsInfo} if invalid.
    *
-   * <p>The cgroups directory is created at most once per Blaze instance.
+   * <p>For v1, we only care about the memory hierarchy.
    *
-   * @return A CgroupsInfo object that defines the cgroups directory that Blaze can use for
-   *     sub-processes. The Blaze process itself is not moved into this directory.
-   * @throws IOException If there are errors reading any of the required files.
+   * @param procMountsFile the /proc/self/mounts file.
    */
-  public static CgroupsInfo createBlazeSpawnsCgroup() throws IOException {
-    return createBlazeSpawnsCgroup(PROC_SELF_MOUNTS_PATH, PROC_SELF_CGROUP_PATH);
-  }
-
   @VisibleForTesting
-  static CgroupsInfo createBlazeSpawnsCgroup(String procSelfMountsPath, String procSelfCgroupPath)
-      throws IOException {
-    Pair<File, Boolean> cgroupsMount = getCgroupMountInfo(new File(procSelfMountsPath));
-    File blazeSpawnsDir;
-    File cgroupsMountPoint = cgroupsMount.first;
-    if (cgroupsMount.second) {
-      File blazeProcessCgroupDir =
-          getBlazeProcessCgroupDir(cgroupsMountPoint, 0, procSelfCgroupPath);
-      // In cgroups v2, we need to step back from the leaf node to make a further hierarchy.
-      blazeSpawnsDir =
-          new File(
-              blazeProcessCgroupDir.getParentFile(),
-              "blaze_" + ProcessHandle.current().pid() + "_spawns.slice");
-      blazeSpawnsDir.mkdirs();
-      blazeSpawnsDir.deleteOnExit();
-      setSubtreeControllers(blazeSpawnsDir);
-      logger.atInfo().log("Creating cgroups v2 node at %s", blazeSpawnsDir);
-      return new CgroupsInfo(true, blazeSpawnsDir, cgroupsMountPoint);
-    } else {
-      int memoryHierarchy = getMemoryHierarchy(new File(procSelfCgroupPath));
-      File blazeProcessCgroupDir =
-          getBlazeProcessCgroupDir(cgroupsMountPoint, memoryHierarchy, procSelfCgroupPath);
-      blazeSpawnsDir =
-          new File(blazeProcessCgroupDir, "blaze_" + ProcessHandle.current().pid() + "_spawns");
-      blazeSpawnsDir.mkdirs();
-      blazeSpawnsDir.deleteOnExit();
-      logger.atInfo().log("Creating cgroups v1 node at %s", blazeSpawnsDir);
-      return new CgroupsInfo(false, blazeSpawnsDir, cgroupsMountPoint);
+  static CgroupsInfo getRootCgroup(File procMountsFile) {
+    if (OS.getCurrent() != OS.LINUX) {
+      return new InvalidCgroupsInfo(
+          Type.ROOT, /* version= */ null, "Croups is not supported on non-linux environments.");
     }
+
+    List<String> procMountsContents;
+    try {
+      procMountsContents = Files.readLines(procMountsFile, UTF_8);
+    } catch (IOException e) {
+      return new InvalidCgroupsInfo(Type.ROOT, /* version= */ null, e);
+    }
+    File v1RootDir = null;
+    File v2RootDir = null;
+    for (String s : procMountsContents) {
+      Matcher m = CGROUPS_MOUNT_PATTERN.matcher(s);
+      if (m.matches()) {
+        if (m.group(1).isEmpty()) {
+          // v1
+          if (m.group(3).contains("memory")) {
+            // For now, we only care about the memory cgroup
+            v1RootDir = new File(m.group(2));
+          }
+        } else {
+          v2RootDir = new File(m.group(2));
+        }
+      }
+    }
+    // If we found the memory controller in v1, we use that, just in case we have a hybrid system
+    // where some controllers are v1 and some are v2. It would be harder to detect if v2 has the
+    // memory controller
+    if (v1RootDir != null) {
+      return new CgroupsInfoV1(Type.ROOT, v1RootDir);
+    }
+    if (v2RootDir != null) {
+      return new CgroupsInfoV2(Type.ROOT, v2RootDir);
+    }
+    return new InvalidCgroupsInfo(
+        Type.ROOT,
+        /* version= */ null,
+        String.format(
+            "No cgroups mounted in %s: %s", procMountsFile.getPath(), procMountsContents));
   }
 
   /**
-   * Sets the subtree controllers we need. This also checks that the controllers are available.
-   *
-   * @param blazeDir A directory in the cgroups hierarchy.
-   * @throws IOException If reading or writing the {@code cgroup.controllers} or {@code
-   *     cgroup.subtree_control} file fails.
-   * @throws IllegalStateException if the {@code memory} and {code pids} controllers are either not
-   *     available or cannot be set for subtrees.
+   * Returns the singleton {@link Type.BLAZE_SPAWNS} cgroup created under the root cgroup, {@link
+   * InvalidCgroupsInfo} if invalid.
    */
-  private static void setSubtreeControllers(File blazeDir) throws IOException {
-    var controllers =
-        Joiner.on(' ').join(Files.readLines(new File(blazeDir, "cgroup.controllers"), UTF_8));
-    if (!(controllers.contains("memory") && controllers.contains("pids"))) {
-      throw new IllegalStateException(
-          String.format(
-              "Required controllers 'memory' and 'pids' not found in %s/cgroup.controllers",
-              blazeDir));
+  public static CgroupsInfo getBlazeSpawnsCgroup() {
+    return blazeSpawnsCgroupSupplier.get();
+  }
+
+  private static CgroupsInfo createBlazeSpawnsCgroup() {
+    if (!rootCgroup.exists()) {
+      return new InvalidCgroupsInfo(
+          Type.BLAZE_SPAWNS, rootCgroup.getVersion(), "Root cgroup does not exist.");
     }
-    var subtreeControllers =
-        Joiner.on(' ').join(Files.readLines(new File(blazeDir, "cgroup.subtree_control"), UTF_8));
-    if (!subtreeControllers.contains("memory") || !subtreeControllers.contains("pids")) {
-      Files.asCharSink(new File(blazeDir, "cgroup.subtree_control"), UTF_8)
-          .write("+memory +pids\n");
+    return rootCgroup.createBlazeSpawnsCgroup(PROC_SELF_CGROUP_PATH);
+  }
+
+  /**
+   * Creates a cgroups directory for Blaze to place spawns in.
+   *
+   * <p>This cgroups directory is created at most once per Blaze instance.
+   *
+   * @param procSelfCgroupPath path to the <code>/proc/self/cgroup</code> file
+   * @return A CgroupsInfo object representing the created cgroup that Blaze can use for
+   *     sub-processes (the Blaze process itself is not moved into this directory). If unable to
+   *     create, returns an {@link InvalidCgroupsInfo} containing the exception.
+   */
+  public abstract CgroupsInfo createBlazeSpawnsCgroup(String procSelfCgroupPath);
+
+  /** The version of Cgroups that is currently being used. */
+  public enum Version {
+    V1,
+    V2,
+  }
+
+  @Nullable protected Version version;
+
+  /**
+   * The types of cgroups relevant to Blaze:
+   *
+   * <ul>
+   *   <li>ROOT: corresponds to the root cgroup where * the hierarchy is mounted at; one of
+   *       "/dev/cgroup/{controller}" or "/sys/fs/cgroup".
+   *   <li>BLAZE_SPAWNS: corresponds the overarching cgroup that contains children {@link
+   *       Type.SPAWN} cgroups.
+   *   <li>SPAWN: corresponds to the cgroup for a single spawn - this could be a locally executed
+   *       action or a worker process.
+   * </ul>
+   */
+  public enum Type {
+    ROOT,
+    BLAZE_SPAWNS,
+    SPAWN,
+  }
+
+  protected Type type;
+
+  /**
+   * This is the directory where the cgroup is in, any related files pertaining to limits / resource
+   * usage or child cgroups (nested directories) are found here.
+   */
+  @Nullable protected final File cgroupDir;
+
+  public CgroupsInfo(Type type, @Nullable Version version, @Nullable File cgroupDir) {
+    this.version = version;
+    this.type = type;
+    this.cgroupDir = cgroupDir;
+    // Valid.
+    if (exists()) {
+      logger.atInfo().log(
+          "Successfully found / created %s (%s) cgroup at %s", version, type, cgroupDir.getPath());
     }
   }
 
-  public boolean isCgroupsV2() {
-    return isCgroupsV2;
+  /** Returns whether the cgroup at {@code cgroupDir} exists. */
+  public boolean exists() {
+    return cgroupDir != null && cgroupDir.exists() && cgroupDir.isDirectory();
+  }
+
+  /** Returns whether Blaze can write to the current cgroup at {@code cgroupDir}. */
+  public boolean canWrite() {
+    return exists() && cgroupDir.canWrite();
   }
 
   /** A cgroups directory for this Blaze instance to put sandboxes in. */
@@ -141,183 +196,90 @@
     return cgroupDir;
   }
 
-  /** The place where the cgroups (memory) file system is mounted. */
-  public File getMountPoint() {
-    return mountPoint;
+  @Nullable
+  public Version getVersion() {
+    return version;
+  }
+
+  public Type getType() {
+    return type;
+  }
+
+  public int getMemoryUsageInKb() {
+    return 0;
+  }
+
+  public int getMemoryUsageInKbFromFile(String filename) {
+    try {
+      String val = Files.readLines(new File(cgroupDir, filename), UTF_8).get(0);
+      return (int) (Long.parseLong(val) / 1024);
+    } catch (IOException e) {
+      return 0;
+    }
+  }
+
+  public void addProcess(long pid) throws IOException {
+    Files.asCharSink(new File(cgroupDir, "cgroup.procs"), UTF_8).write(Long.toString(pid));
   }
 
   /**
-   * Reads from the given file (e.g. /proc/mounts) where cgroups are mounted. If both cgroups v1 and
-   * cgroups v2 are mounted, the one that has the memory controller is used.
+   * Creates a cgroups directory for individual spawns (local / workers).
    *
-   * <p>In cgroups v1, a typical mount line looks like this (note {@code memory} in the options):
+   * <p>Has to be called from a {@link Type.BLAZE_SPAWNS} cgroup.
    *
-   * <pre>
-   * cgroup /dev/cgroup/memory cgroup rw,memory,hugetlb 0 0
-   * </pre>
-   *
-   * In cgroups v2, there is only one relevant line, and it looks like this:
-   *
-   * <pre>
-   * cgroup2 /sys/fs/cgroup cgroup2 rw,[...] 0 0
-   * </pre>
-   *
-   * @param procMountsPath Paths of the mounts file, e.g. /proc/mounts
-   * @return Pair of
-   *     <ol>
-   *       <li>the path of the cgroups mount (for cgroups v1, this is the memory hierarchy) and
-   *       <li>whether this is cgroups v2.
-   *     </ol>
-   *
-   * @throws IOException If there are errors reading the given file.
+   * @param dirName the directory name of the spawn's cgroup.
+   * @param memoryLimitMb memory limit in Mb to set on the cgroup. If 0, no limit is set.
+   * @return an instance of the spawn's cgroup; if unable to create, returns an {@link
+   *     InvalidCgroupsInfo} containing the exception.
    */
-  @VisibleForTesting
-  static Pair<File, Boolean> getCgroupMountInfo(File procMountsPath) throws IOException {
-    var procMountContents = Files.readLines(procMountsPath, UTF_8);
-    Pair<File, Boolean> v1 = null;
-    Pair<File, Boolean> v2 = null;
-    for (String s : procMountContents) {
-      Matcher m = CGROUPS_MOUNT_PATTERN.matcher(s);
-      if (m.matches()) {
-        if (m.group(1).isEmpty()) {
-          // v1
-          if (m.group(3).contains("memory")) {
-            // For now, we only care about the memory cgroup
-            v1 = Pair.of(new File(m.group(2)), false);
-          }
-        } else {
-          // v2
-          v2 = Pair.of(new File(m.group(2)), true);
-        }
-      }
-    }
-    // If we found the memory controller in v1, we use that, just in case we have a hybrid system
-    // where some controllers are v1 and some are v2. It would be harder to detect if v2 has the
-    // memory controller
-    if (v1 != null) {
-      return v1;
-    }
-    if (v2 != null) {
-      return v2;
-    }
-    throw new IllegalStateException(
-        "Cgroups requested, but no applicable cgroups are mounted on this machine");
-  }
+  public abstract CgroupsInfo createIndividualSpawnCgroup(String dirName, int memoryLimitMb);
 
   /**
-   * Returns the number of the memory cgroups v1 hierarchy.
-   *
-   * <p>The <code>/proc/self/cgroup</code> file look like this in v1:
-   *
-   * <pre>
-   * 8:net:/some/path
-   * 7:memory,hugetlb:/some/other/path
-   * ...
-   * </pre>
-   *
-   * In v2, there is only one entry, and it looks something like
-   *
-   * <pre>
-   * 0::/user.slice/user-123.slice/session-1.scope
-   * </pre>
-   *
-   * @param procSelfCgroupPath Path for the <code>/proc/self/cgroup</code> file.
-   * @return The hierarchy number for the cgroups v1 hierarchy that contains the memory controller.
-   * @throws IOException If there are errors reading the file.
+   * Represents an invalid cgroup so that we can distinguish between whether a cgroup was not meant
+   * to be created (null) or if it was attempted but failed.
    */
-  @VisibleForTesting
-  static int getMemoryHierarchy(File procSelfCgroupPath) throws IOException {
-    List<String> devCgroupContents = Files.readLines(procSelfCgroupPath, UTF_8);
-    for (String s : devCgroupContents) {
-      if (s.contains("memory")) {
-        return Integer.parseInt(Splitter.on(":").split(s).iterator().next());
-      }
-    }
-    throw new IllegalStateException(
-        String.format(
-            "Cgroups v1 requested, but no memory cgroup found in %s", procSelfCgroupPath));
-  }
+  public static class InvalidCgroupsInfo extends CgroupsInfo {
 
-  /**
-   * Returns the path of the memory cgroups node that Blaze itself runs inside.
-   *
-   * @param mountPoint Where the cgroups hierarchy (with the memory controller for v1) is mounted.
-   * @param memoryHierarchyId The v1 hierarchy that contains the memory controller, or 0 for v2.
-   * @param procSelfPath The path of the <code>/proc/self/cgroup</code> file.
-   * @return A <code>File</code> object of the cgroup directory of the current process.
-   * @throws IOException If the given file cannot be read.
-   */
-  @VisibleForTesting
-  static File getBlazeProcessCgroupDir(File mountPoint, int memoryHierarchyId, String procSelfPath)
-      throws IOException {
-    var procSelfCgroupContents = Files.readLines(new File(procSelfPath), UTF_8);
-    if (procSelfCgroupContents.isEmpty()) {
-      throw new IOException("Cgroups requested, but /proc/self/cgroup is empty");
-    }
-    File cgroupsNode = null;
-    for (String s : procSelfCgroupContents) {
-      List<String> parts = Splitter.on(":").limit(3).splitToList(s);
-      if (parts.size() == 3 && Integer.parseInt(parts.get(0)) == memoryHierarchyId) {
-        String path = parts.get(2);
-        if (path.startsWith(File.pathSeparator)) {
-          path = path.substring(1);
-        }
-        cgroupsNode = new File(mountPoint, path);
-        break;
-      }
-    }
-    if (cgroupsNode == null) {
-      throw new IllegalStateException("Found no memory cgroups entries in '" + procSelfPath + "'");
-    }
-    if (!cgroupsNode.exists()) {
-      throw new IllegalStateException("Cgroups node '" + cgroupsNode + "' does not exist");
-    }
-    if (!cgroupsNode.isDirectory()) {
-      throw new IllegalStateException("Cgroups node " + cgroupsNode + " is not a directory");
-    }
-    if (!cgroupsNode.canWrite()) {
-      throw new IllegalStateException("Cgroups node " + cgroupsNode + " is not writable");
-    }
-    return cgroupsNode;
-  }
+    private final Exception exception;
 
-  public static CgroupsInfo getBlazeSpawnsCgroup() throws IOException {
-    if (instance == null) {
-      synchronized (CgroupsInfo.class) {
-        if (instance == null) {
-          instance = createBlazeSpawnsCgroup();
-        }
-      }
+    public InvalidCgroupsInfo(Type type, @Nullable Version version, String errorMessage) {
+      super(type, version, null);
+      this.exception = new IllegalStateException(errorMessage);
+      logger.atInfo().withCause(exception).log("Unable to create cgroup.");
     }
-    return instance;
-  }
 
-  /**
-   * Creates a cgroups directory with the given memory limit.
-   *
-   * @param memoryLimit Memory limit in megabytes (MiB).
-   * @param dirName Base name of the directory created. In cgroups v2, <code>.scope</code> gets
-   *     appended.
-   */
-  public static CgroupsInfo createMemoryLimitCgroupDir(
-      CgroupsInfo blazeSpawnsCgroup, String dirName, int memoryLimit) throws IOException {
-    File cgroupsDir;
-    if (blazeSpawnsCgroup.isCgroupsV2()) {
-      cgroupsDir = new File(blazeSpawnsCgroup.getCgroupDir(), dirName + ".scope");
-      cgroupsDir.mkdirs();
-      cgroupsDir.deleteOnExit();
-      // In cgroups v2, we need to propagate the controllers into new subdirs.
-      Files.asCharSink(new File(cgroupsDir, "memory.oom.group"), UTF_8).write("1\n");
-      Files.asCharSink(new File(cgroupsDir, "memory.max"), UTF_8)
-          .write(Long.toString(memoryLimit * 1024L * 1024L));
-    } else {
-      cgroupsDir = new File(blazeSpawnsCgroup.getCgroupDir(), dirName);
-      cgroupsDir.mkdirs();
-      cgroupsDir.deleteOnExit();
-      Files.asCharSink(new File(cgroupsDir, "memory.limit_in_bytes"), UTF_8)
-          .write(Long.toString(memoryLimit * 1024L * 1024L));
+    public InvalidCgroupsInfo(Type type, @Nullable Version version, Exception exception) {
+      super(type, version, null);
+      logger.atInfo().withCause(exception).log("Unable to create cgroup.");
+      this.exception = exception;
     }
-    return new CgroupsInfo(
-        blazeSpawnsCgroup.isCgroupsV2(), cgroupsDir, blazeSpawnsCgroup.getMountPoint());
+
+    @Override
+    public boolean exists() {
+      return false;
+    }
+
+    @Override
+    public boolean canWrite() {
+      return false;
+    }
+
+    public Exception getException() {
+      return exception;
+    }
+
+    @Override
+    public CgroupsInfo createBlazeSpawnsCgroup(String procSelfCgroupPath) {
+      return new InvalidCgroupsInfo(
+          Type.BLAZE_SPAWNS,
+          getVersion(),
+          "Unable to create BLAZE_SPAWNS cgroup from an invalid cgroup.");
+    }
+
+    @Override
+    public CgroupsInfo createIndividualSpawnCgroup(String dirName, int memoryLimitMb) {
+      return new InvalidCgroupsInfo(
+          Type.SPAWN, getVersion(), "Unable to create SPAWN cgroup from an invalid cgroup.");
+    }
   }
 }
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfoV1.java b/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfoV1.java
new file mode 100644
index 0000000..7df0f57
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfoV1.java
@@ -0,0 +1,109 @@
+// Copyright 2024 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.sandbox;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.common.base.Splitter;
+import com.google.common.io.Files;
+import java.io.File;
+import java.io.IOException;
+import java.util.List;
+import javax.annotation.Nullable;
+
+/** Represents a v1 cgroup. */
+public class CgroupsInfoV1 extends CgroupsInfo {
+  public CgroupsInfoV1(Type type, @Nullable File cgroupDir) {
+    super(type, Version.V1, cgroupDir);
+  }
+
+  @Override
+  public CgroupsInfo createBlazeSpawnsCgroup(String procSelfCgroupPath) {
+    checkArgument(
+        type == Type.ROOT, "Should only be creating the Blaze spawns cgroup from the root cgroup.");
+    File blazeProcessCgroupDir;
+    try {
+      blazeProcessCgroupDir = getBlazeProcessCgroupDir(cgroupDir, procSelfCgroupPath);
+    } catch (Exception e) {
+      return new InvalidCgroupsInfo(Type.BLAZE_SPAWNS, getVersion(), e);
+    }
+    File blazeSpawnsDir =
+        new File(blazeProcessCgroupDir, "blaze_" + ProcessHandle.current().pid() + "_spawns");
+    blazeSpawnsDir.mkdirs();
+    blazeSpawnsDir.deleteOnExit();
+    return new CgroupsInfoV1(Type.BLAZE_SPAWNS, blazeSpawnsDir);
+  }
+
+  @Override
+  public CgroupsInfo createIndividualSpawnCgroup(String dirName, int memoryLimitMb) {
+    checkArgument(
+        type == Type.BLAZE_SPAWNS,
+        "Should only be creating the individual spawn's cgroup from the Blaze spawns cgroup.");
+    if (!canWrite()) {
+      return new InvalidCgroupsInfo(
+          Type.SPAWN,
+          getVersion(),
+          String.format("Cgroup %s is invalid, unable to create spawn's cgroup here.", cgroupDir));
+    }
+    File spawnCgroupDir = new File(cgroupDir, dirName);
+    spawnCgroupDir.mkdirs();
+    spawnCgroupDir.deleteOnExit();
+    try {
+      if (memoryLimitMb > 0) {
+        Files.asCharSink(new File(spawnCgroupDir, "memory.limit_in_bytes"), UTF_8)
+            .write(Long.toString(memoryLimitMb * 1024L * 1024L));
+      }
+    } catch (Exception e) {
+      return new InvalidCgroupsInfo(Type.SPAWN, getVersion(), e);
+    }
+    return new CgroupsInfoV1(Type.SPAWN, spawnCgroupDir);
+  }
+
+  /**
+   * Returns the path to the cgroup containing the Blaze process.
+   *
+   * <p>The <code>/proc/self/cgroup</code> file look like this in v1:
+   *
+   * <pre>
+   * 8:net:/some/path
+   * 7:memory,hugetlb:/some/other/path
+   * ...
+   * </pre>
+   *
+   * @param mountPoint the directory where the cgroup hierarchy is mounted.
+   * @param procSelfCgroupPath path for the /proc/self/cgroup file.
+   * @throws IOException if there are errors reading the given procs cgroup file.
+   */
+  private static File getBlazeProcessCgroupDir(File mountPoint, String procSelfCgroupPath)
+      throws IOException {
+    List<String> controllers = Files.readLines(new File(procSelfCgroupPath), UTF_8);
+    String memoryController =
+        controllers.stream()
+            .filter(controller -> controller.contains("memory"))
+            .findFirst()
+            .orElseThrow(
+                () ->
+                    new IllegalStateException(
+                        "Found no memory cgroup entries in '" + procSelfCgroupPath + "'"));
+    List<String> parts = Splitter.on(":").limit(3).splitToList(memoryController);
+    return new File(mountPoint, parts.get(2));
+  }
+
+  @Override
+  public int getMemoryUsageInKb() {
+    return getMemoryUsageInKbFromFile("memory.usage_in_bytes");
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfoV2.java b/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfoV2.java
new file mode 100644
index 0000000..1f62f8d
--- /dev/null
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/CgroupsInfoV2.java
@@ -0,0 +1,139 @@
+// Copyright 2024 The Bazel Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package com.google.devtools.build.lib.sandbox;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.nio.charset.StandardCharsets.UTF_8;
+
+import com.google.common.base.Joiner;
+import com.google.common.base.Splitter;
+import com.google.common.io.Files;
+import java.io.File;
+import java.io.IOException;
+import java.util.List;
+import javax.annotation.Nullable;
+
+/** Represents a v2 cgroup. */
+public class CgroupsInfoV2 extends CgroupsInfo {
+
+  public CgroupsInfoV2(Type type, @Nullable File cgroupDir) {
+    super(type, Version.V2, cgroupDir);
+  }
+
+  @Override
+  public CgroupsInfo createBlazeSpawnsCgroup(String procSelfCgroupPath) {
+    checkArgument(
+        type == Type.ROOT, "Should only be creating the Blaze spawns cgroup from the root cgroup.");
+    File blazeProcessCgroupDir;
+    File blazeSpawnsDir;
+    try {
+      blazeProcessCgroupDir = getBlazeProcessCgroupDir(cgroupDir, procSelfCgroupPath);
+      // In cgroups v2, we need to step back from the leaf node to make a further hierarchy.
+      blazeSpawnsDir =
+          new File(
+              blazeProcessCgroupDir.getParentFile(),
+              "blaze_" + ProcessHandle.current().pid() + "_spawns.slice");
+      blazeSpawnsDir.mkdirs();
+      blazeSpawnsDir.deleteOnExit();
+      setSubtreeControllers(blazeSpawnsDir);
+    } catch (Exception e) {
+      return new InvalidCgroupsInfo(Type.BLAZE_SPAWNS, getVersion(), e);
+    }
+
+    return new CgroupsInfoV2(Type.BLAZE_SPAWNS, blazeSpawnsDir);
+  }
+
+  @Override
+  public CgroupsInfo createIndividualSpawnCgroup(String dirName, int memoryLimitMb) {
+    checkArgument(
+        type == Type.BLAZE_SPAWNS,
+        "Should only be creating the individual spawn's cgroup from the Blaze spawns cgroup.");
+    if (!canWrite()) {
+      return new InvalidCgroupsInfo(
+          Type.SPAWN,
+          getVersion(),
+          String.format("Cgroup %s is invalid, unable to create spawn's cgroup here.", cgroupDir));
+    }
+    File spawnCgroupDir = new File(cgroupDir, dirName + ".scope");
+    spawnCgroupDir.mkdirs();
+    spawnCgroupDir.deleteOnExit();
+    try {
+      if (memoryLimitMb > 0) {
+        // In cgroups v2, we need to propagate the controllers into new subdirs.
+        Files.asCharSink(new File(spawnCgroupDir, "memory.oom.group"), UTF_8).write("1\n");
+        Files.asCharSink(new File(spawnCgroupDir, "memory.max"), UTF_8)
+            .write(Long.toString(memoryLimitMb * 1024L * 1024L));
+      }
+    } catch (Exception e) {
+      return new InvalidCgroupsInfo(Type.SPAWN, getVersion(), e);
+    }
+    return new CgroupsInfoV2(Type.SPAWN, spawnCgroupDir);
+  }
+
+  @Override
+  public int getMemoryUsageInKb() {
+    return getMemoryUsageInKbFromFile("memory.current");
+  }
+
+  /**
+   * Returns the path to the cgroup containing the Blaze process.
+   *
+   * <p>In v2, there is only one entry, and it looks something like this:
+   *
+   * <pre>
+   * 0::/user.slice/user-123.slice/session-1.scope
+   * </pre>
+   *
+   * @param mountPoint the directory where the cgroup hierarchy is mounted.
+   * @param procSelfCgroupPath path for the /proc/self/cgroup file.
+   * @throws IOException if there are errors reading the given procs cgroup file.
+   */
+  private static File getBlazeProcessCgroupDir(File mountPoint, String procSelfCgroupPath)
+      throws IOException {
+    List<String> contents = Files.readLines(new File(procSelfCgroupPath), UTF_8);
+    if (contents.isEmpty()) {
+      throw new IllegalStateException(
+          "Found no memory cgroup entries in '" + procSelfCgroupPath + "'");
+    }
+    List<String> parts = Splitter.on(":").limit(3).splitToList(contents.get(0));
+    return new File(mountPoint, parts.get(2));
+  }
+
+  /**
+   * Sets the subtree controllers we need. This also checks that the controllers are available.
+   *
+   * @param blazeDir A directory in the cgroups hierarchy.
+   * @throws IOException If reading or writing the {@code cgroup.controllers} or {@code
+   *     cgroup.subtree_control} file fails.
+   * @throws IllegalStateException if the {@code memory} and {code pids} controllers are either not
+   *     available or cannot be set for subtrees.
+   */
+  private static void setSubtreeControllers(File blazeDir) throws IOException {
+    var controllers =
+        Joiner.on(' ').join(Files.readLines(new File(blazeDir, "cgroup.controllers"), UTF_8));
+    if (!(controllers.contains("memory") && controllers.contains("pids"))) {
+      throw new IllegalStateException(
+          String.format(
+              "Required controllers 'memory' and 'pids' not found in %s/cgroup.controllers",
+              blazeDir));
+    }
+    var subtreeControllers =
+        Joiner.on(' ').join(Files.readLines(new File(blazeDir, "cgroup.subtree_control"), UTF_8));
+    if (!subtreeControllers.contains("memory") || !subtreeControllers.contains("pids")) {
+      Files.asCharSink(new File(blazeDir, "cgroup.subtree_control"), UTF_8)
+          .write("+memory +pids\n");
+    }
+  }
+}
diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/LinuxSandboxedSpawnRunner.java b/src/main/java/com/google/devtools/build/lib/sandbox/LinuxSandboxedSpawnRunner.java
index 4a7327d..f4079d9 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/LinuxSandboxedSpawnRunner.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/LinuxSandboxedSpawnRunner.java
@@ -326,10 +326,9 @@
       // We put the sandbox inside a unique subdirectory using the context's ID. This ID is
       // unique per spawn run by this spawn runner.
       CgroupsInfo sandboxCgroup =
-          CgroupsInfo.createMemoryLimitCgroupDir(
-              CgroupsInfo.getBlazeSpawnsCgroup(),
-              "sandbox_" + context.getId(),
-              sandboxOptions.memoryLimitMb);
+          CgroupsInfo.getBlazeSpawnsCgroup()
+              .createIndividualSpawnCgroup(
+                  "sandbox_" + context.getId(), sandboxOptions.memoryLimitMb);
       cgroupsDir = sandboxCgroup.getCgroupDir().toString();
       commandLineBuilder.setCgroupsDir(cgroupsDir);
     }
diff --git a/src/main/java/com/google/devtools/build/lib/worker/SandboxedWorker.java b/src/main/java/com/google/devtools/build/lib/worker/SandboxedWorker.java
index afa3236..cc056fe 100644
--- a/src/main/java/com/google/devtools/build/lib/worker/SandboxedWorker.java
+++ b/src/main/java/com/google/devtools/build/lib/worker/SandboxedWorker.java
@@ -198,10 +198,9 @@
       if (hardenedSandboxOptions.memoryLimit() > 0) {
         // We put the sandbox inside a unique subdirectory using the worker's ID.
         CgroupsInfo workerCgroup =
-            CgroupsInfo.createMemoryLimitCgroupDir(
-                CgroupsInfo.getBlazeSpawnsCgroup(),
-                "worker_sandbox_" + workerId,
-                hardenedSandboxOptions.memoryLimit());
+            CgroupsInfo.getBlazeSpawnsCgroup()
+                .createIndividualSpawnCgroup(
+                    "worker_sandbox_" + workerId, hardenedSandboxOptions.memoryLimit());
         cgroupsDir = workerCgroup.getCgroupDir().toString();
         commandLineBuilder.setCgroupsDir(cgroupsDir);
       }
diff --git a/src/test/java/com/google/devtools/build/lib/sandbox/CgroupsInfoTest.java b/src/test/java/com/google/devtools/build/lib/sandbox/CgroupsInfoTest.java
index 6d21a73..95cfa06 100644
--- a/src/test/java/com/google/devtools/build/lib/sandbox/CgroupsInfoTest.java
+++ b/src/test/java/com/google/devtools/build/lib/sandbox/CgroupsInfoTest.java
@@ -14,22 +14,23 @@
 package com.google.devtools.build.lib.sandbox;
 
 import static com.google.common.truth.Truth.assertThat;
+import static java.nio.charset.StandardCharsets.UTF_8;
 
 import com.google.common.io.Files;
-import com.google.devtools.build.lib.util.Pair;
+import com.google.devtools.build.lib.sandbox.CgroupsInfo.InvalidCgroupsInfo;
 import com.google.devtools.build.lib.vfs.util.FsApparatus;
 import java.io.File;
 import java.io.IOException;
-import java.nio.charset.StandardCharsets;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 
-/** Tests for {@link CgroupsInfo}. */
+/** Tests for {@link CgroupsInfo}, {@link CgroupsInfoV1}, {@link CgroupsInfoV2}. */
 @RunWith(JUnit4.class)
 public class CgroupsInfoTest {
   private final FsApparatus scratch = FsApparatus.newNative();
+
   /** We use this pseudo-root to get around not being able to replace absolute paths. */
   private String root;
 
@@ -39,7 +40,7 @@
   }
 
   @Test
-  public void testGetCgroupMountInfo_v1() throws IOException {
+  public void testGetRootCgroup_v1() throws IOException {
     String pathString =
         createFakeAbsoluteFile(
             "/proc/self/mounts",
@@ -49,26 +50,32 @@
             "cgroup /dev/cgroup/job cgroup rw,job 0 0",
             "cgroup /dev/cgroup/memory cgroup rw,memory,hugetlb 0 0",
             "proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0");
-    Pair<File, Boolean> cgroupsMountInfo = CgroupsInfo.getCgroupMountInfo(new File(pathString));
-    assertThat(cgroupsMountInfo.second).isFalse();
-    assertThat(cgroupsMountInfo.first.toString()).isEqualTo("/dev/cgroup/memory");
+
+    CgroupsInfo cgroup = CgroupsInfo.getRootCgroup(new File(pathString));
+
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.ROOT);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V1);
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo("/dev/cgroup/memory");
   }
 
   @Test
-  public void testGetCgroupMountInfo_v2() throws IOException {
+  public void testGetRootCgroup_v2() throws IOException {
     String pathString =
         createFakeAbsoluteFile(
             "/proc/self/mounts",
             "sysfs /sys sysfs rw,nosuid,nodev,noexec,relatime 0 0",
             "cgroup2 /sys/fs/cgroup cgroup2 rw,memory_recursiveprot 0 0",
             "proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0");
-    Pair<File, Boolean> cgroupsMountInfo = CgroupsInfo.getCgroupMountInfo(new File(pathString));
-    assertThat(cgroupsMountInfo.second).isTrue();
-    assertThat(cgroupsMountInfo.first.toString()).isEqualTo("/sys/fs/cgroup");
+
+    CgroupsInfo cgroup = CgroupsInfo.getRootCgroup(new File(pathString));
+
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.ROOT);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo("/sys/fs/cgroup");
   }
 
   @Test
-  public void testGetCgroupMountInfo_mixed_v1_has_memory() throws IOException {
+  public void testGetRootCgroup_mixed_v1_has_memory() throws IOException {
     String pathString =
         createFakeAbsoluteFile(
             "/proc/self/mounts",
@@ -77,13 +84,16 @@
             "cgroup /dev/cgroup/job cgroup rw,job 0 0",
             "cgroup /dev/cgroup/memory cgroup rw,memory,hugetlb 0 0",
             "proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0");
-    Pair<File, Boolean> cgroupsMountInfo = CgroupsInfo.getCgroupMountInfo(new File(pathString));
-    assertThat(cgroupsMountInfo.second).isFalse();
-    assertThat(cgroupsMountInfo.first.toString()).isEqualTo("/dev/cgroup/memory");
+
+    CgroupsInfo cgroup = CgroupsInfo.getRootCgroup(new File(pathString));
+
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.ROOT);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V1);
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo("/dev/cgroup/memory");
   }
 
   @Test
-  public void testGetCgroupMountInfo_mixed_v2_has_memory() throws IOException {
+  public void testGetRootCgroup_mixed_v2_has_memory() throws IOException {
     String pathString =
         createFakeAbsoluteFile(
             "/proc/self/mount",
@@ -92,28 +102,20 @@
             "cgroup /dev/cgroup/job cgroup rw,job 0 0",
             "cgroup /dev/cgroup/io cgroup rw,io 0 0",
             "proc /proc proc rw,nosuid,nodev,noexec,relatime 0 0");
-    Pair<File, Boolean> cgroupsMountInfo = CgroupsInfo.getCgroupMountInfo(new File(pathString));
-    assertThat(cgroupsMountInfo.second).isTrue();
-    assertThat(cgroupsMountInfo.first.toString()).isEqualTo("/sys/fs/cgroup");
+
+    CgroupsInfo cgroup = CgroupsInfo.getRootCgroup(new File(pathString));
+
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.ROOT);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo("/sys/fs/cgroup");
   }
 
   @Test
-  public void testGetCgroupsHierarchy_v1() throws IOException {
-    String pathString =
-        createFakeAbsoluteFile(
-            "/proc/self/cgroup",
-            "8:net:/netdir/action-6",
-            "7:memory,hugetlb:/memdir/action-6",
-            "6:job:/jobdir/action-16",
-            "5:io:/iodir/action-1");
-    int hierarchy = CgroupsInfo.getMemoryHierarchy(new File(pathString));
-    assertThat(hierarchy).isEqualTo(7);
-  }
-
-  @Test
-  public void testGetCgroupsNode_v1() throws IOException {
+  public void testCreateBlazeSpawnsCgroup_v1() throws IOException {
     String mountPath = root + "/dev/cgroup/memory";
-    String pathString =
+    CgroupsInfo rootCgroup =
+        new CgroupsInfoV1(CgroupsInfo.Type.ROOT, /* cgroupDir= */ new File(mountPath));
+    String procSelfCgroupPath =
         createFakeAbsoluteFile(
             "/proc/self/cgroup",
             "8:net:/netdir/action-6",
@@ -121,89 +123,219 @@
             "6:job:/jobdir/action-16",
             "5:io:/iodir/action-1");
     scratch.dir(root + "/dev/cgroup/memory/memdir/action-6").createDirectoryAndParents();
-    File cgroupsMountInfo =
-        CgroupsInfo.getBlazeProcessCgroupDir(new File(mountPath), 7, pathString);
-    assertThat(cgroupsMountInfo.getAbsolutePath()).isEqualTo(mountPath + "/memdir/action-6");
+    String blazeSpawnsPath =
+        root
+            + "/dev/cgroup/memory/memdir/action-6/blaze_"
+            + ProcessHandle.current().pid()
+            + "_spawns";
+    scratch.dir(blazeSpawnsPath).createDirectoryAndParents();
+
+    CgroupsInfo cgroup = rootCgroup.createBlazeSpawnsCgroup(procSelfCgroupPath);
+
+    assertThat(cgroup.getCgroupDir().exists()).isTrue();
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo(blazeSpawnsPath);
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.BLAZE_SPAWNS);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V1);
   }
 
   @Test
-  public void testGetCgroupsNode_v2() throws IOException {
+  public void testCreateBlazeSpawnsCgroup_v2() throws IOException {
     String mountPath = root + "/sys/fs/cgroup";
-    String procSelfCgroup =
+    CgroupsInfo rootCgroup =
+        new CgroupsInfoV2(CgroupsInfo.Type.ROOT, /* cgroupDir= */ new File(mountPath));
+    String procSelfCgroupPath =
         createFakeAbsoluteFile("/proc/self/cgroup", "0::/user.slice/session.scope");
-    scratch.dir(mountPath + "/user.slice/session.scope").createDirectoryAndParents();
-    File cgroupsMountInfo =
-        CgroupsInfo.getBlazeProcessCgroupDir(new File(mountPath), 0, procSelfCgroup);
-    assertThat(cgroupsMountInfo.getAbsolutePath())
-        .isEqualTo(mountPath + "/user.slice/session.scope");
-  }
-
-  @Test
-  public void testCreate_v1() throws IOException {
-    // We actually use the paths from /proc/mount here, so they must include the fake root.
-    String procMount =
-        createFakeAbsoluteFile(
-            "/proc/self/mounts",
-            "sysfs " + root + "/sys sysfs rw,nosuid,nodev,noexec,relatime 0 0",
-            "cgroup " + root + "/dev/cgroup/cpu cgroup rw,cpu,cpuacct 0 0",
-            "cgroup " + root + "/dev/cgroup/io cgroup rw,io 0 0",
-            "cgroup " + root + "/dev/cgroup/job cgroup rw,job 0 0",
-            "cgroup " + root + "/dev/cgroup/memory cgroup rw,memory,hugetlb 0 0",
-            "proc " + root + "/proc proc rw,nosuid,nodev,noexec,relatime 0 0");
-    String procSelfCgroup =
-        createFakeAbsoluteFile(
-            "/proc/self/cgroup",
-            "8:net:/netdir/action-6",
-            "7:memory,hugetlb:/memdir/action-6",
-            "6:job:/jobdir/action-16",
-            "5:io:/iodir/action-1");
-    scratch.dir(root + "/dev/cgroup/memory/memdir/action-6").createDirectoryAndParents();
-    CgroupsInfo cgroupsInfo = CgroupsInfo.createBlazeSpawnsCgroup(procMount, procSelfCgroup);
-    assertThat(cgroupsInfo.isCgroupsV2()).isFalse();
-    assertThat(cgroupsInfo.getCgroupDir().getAbsolutePath())
-        .matches(root + "/dev/cgroup/memory/memdir/action-6/blaze_\\d+_spawns");
-    assertThat(cgroupsInfo.getCgroupDir().exists()).isTrue();
-  }
-
-  @Test
-  public void testCreate_v2() throws IOException {
-    String cgroupsRoot = root + "/sys/fs/cgroup";
-    // We actually use the paths from /proc/mount here, so they must include the fake root.
-    String procMount =
-        createFakeAbsoluteFile(
-            "/proc/mount",
-            "sysfs " + root + "/sys sysfs rw,nosuid,nodev,noexec,relatime 0 0",
-            "cgroup2 " + cgroupsRoot + " cgroup2 rw,memory_recursiveprot 0 0",
-            "proc " + root + "/proc proc rw,nosuid,nodev,noexec,relatime 0 0");
-    String procSelfCgroup =
-        createFakeAbsoluteFile("/proc/self/cgroup", "0::/user.slice/session.scope");
-
-    // We base the new subtree off the parent of the current scope.
-    String userSlice = root + "/sys/fs/cgroup/user.slice";
-    // Even though we create a separate directory off `user.slice`, we must check that the
-    // scope for the current process is writable.
-    scratch.dir(userSlice + "/session.scope").createDirectoryAndParents();
-
-    String blazeSlice = userSlice + "/blaze_" + ProcessHandle.current().pid() + "_spawns.slice";
-    scratch.dir(blazeSlice).createDirectoryAndParents();
-    scratch.file(blazeSlice + "/cgroup.controllers", "memory pids");
+    // In v2, the blaze spawns cgroup is created one step up from where the blaze process is defined
+    // in the /proc/self/cgroup file (defined above). Specifically, here it is in the
+    // ".../user.slice".
+    String blazeSpawnsPath =
+        mountPath + "/user.slice/blaze_" + ProcessHandle.current().pid() + "_spawns.slice";
+    // Even though the blaze spawn's cgroup directory is meant to be created in the method call,
+    // we create it here so that we can prepare the controller files that are expected beforehand.
+    scratch.dir(blazeSpawnsPath).createDirectoryAndParents();
     // Since this controllers file is missing `pids`, we expect that to be written to it.
-    scratch.file(blazeSlice + "/cgroup.subtree_control", "memory");
+    scratch.file(blazeSpawnsPath + "/cgroup.controllers", "memory pids");
+    scratch.file(blazeSpawnsPath + "/cgroup.subtree_control", "memory");
 
-    CgroupsInfo cgroupsInfo = CgroupsInfo.createBlazeSpawnsCgroup(procMount, procSelfCgroup);
+    CgroupsInfo cgroup = rootCgroup.createBlazeSpawnsCgroup(procSelfCgroupPath);
 
-    assertThat(cgroupsInfo.isCgroupsV2()).isTrue();
-    assertThat(cgroupsInfo.getCgroupDir().getAbsolutePath())
-        .matches(root + "/sys/fs/cgroup/user.slice/blaze_\\d+_spawns.slice");
-
-    // This is not what an actual cgroupsv2 file would contain, but it's what we expect to write to
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    assertThat(cgroup.getCgroupDir().exists()).isTrue();
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo(blazeSpawnsPath);
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.BLAZE_SPAWNS);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    // This is not what an actual cgroups v2 file would contain, but it's what we expect to write to
     // it to enable subtree control.
-    assertThat(
-            Files.readLines(
-                new File(blazeSlice + "/cgroup.subtree_control"), StandardCharsets.UTF_8))
+    assertThat(Files.readLines(new File(blazeSpawnsPath + "/cgroup.subtree_control"), UTF_8))
         .containsExactly("+memory +pids");
   }
 
+  @Test
+  public void testCreateIndividualSpawnCgroup_withLimit_v1() throws IOException {
+    String blazeSpawnsPath = root + "/dev/cgroup/memory/memdir/action-6/blaze_1234_spawns";
+    scratch.dir(blazeSpawnsPath).createDirectoryAndParents();
+    CgroupsInfo blazeSpawnsCgroup =
+        new CgroupsInfoV1(CgroupsInfo.Type.BLAZE_SPAWNS, new File(blazeSpawnsPath));
+
+    CgroupsInfo cgroup = blazeSpawnsCgroup.createIndividualSpawnCgroup("spawn_1", 100);
+
+    assertThat(cgroup.exists()).isTrue();
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo(blazeSpawnsPath + "/spawn_1");
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.SPAWN);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V1);
+    assertThat(Files.readLines(new File(blazeSpawnsPath + "/spawn_1/memory.limit_in_bytes"), UTF_8))
+        .containsExactly("104857600");
+  }
+
+  @Test
+  public void testCreateIndividualSpawnCgroup_noLimit_v1() throws IOException {
+    String blazeSpawnsPath = root + "/dev/cgroup/memory/memdir/action-6/blaze_1234_spawns";
+    scratch.dir(blazeSpawnsPath).createDirectoryAndParents();
+    CgroupsInfo blazeSpawnsCgroup =
+        new CgroupsInfoV1(CgroupsInfo.Type.BLAZE_SPAWNS, new File(blazeSpawnsPath));
+
+    CgroupsInfo cgroup = blazeSpawnsCgroup.createIndividualSpawnCgroup("spawn_1", 0);
+
+    assertThat(cgroup.exists()).isTrue();
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo(blazeSpawnsPath + "/spawn_1");
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.SPAWN);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V1);
+    // In reality, cgroups should still create this file automatically, but since we don't have
+    // that in our tests, the memory limits file should not have been created since it isn't written
+    // to.
+    assertThat(new File(blazeSpawnsPath + "/spawn_1/memory.limit_in_bytes").exists()).isFalse();
+  }
+
+  @Test
+  public void testCreateIndividualSpawnCgroup_withLimit_v2() throws IOException {
+    String blazeSpawnsPath = root + "/sys/fs/cgroup/user.slice/blaze_1234_spawns.slice";
+    scratch.dir(blazeSpawnsPath).createDirectoryAndParents();
+    CgroupsInfo blazeSpawnsCgroup =
+        new CgroupsInfoV2(CgroupsInfo.Type.BLAZE_SPAWNS, new File(blazeSpawnsPath));
+
+    CgroupsInfo cgroup = blazeSpawnsCgroup.createIndividualSpawnCgroup("spawn_1", 100);
+
+    assertThat(cgroup.exists()).isTrue();
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo(blazeSpawnsPath + "/spawn_1.scope");
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.SPAWN);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    assertThat(
+            Files.readLines(new File(blazeSpawnsPath + "/spawn_1.scope/memory.oom.group"), UTF_8))
+        .containsExactly("1");
+    assertThat(Files.readLines(new File(blazeSpawnsPath + "/spawn_1.scope/memory.max"), UTF_8))
+        .containsExactly("104857600");
+  }
+
+  @Test
+  public void testCreateIndividualSpawnCgroup_noLimit_v2() throws IOException {
+    String blazeSpawnsPath = root + "/sys/fs/cgroup/user.slice/blaze_1234_spawns.slice";
+    scratch.dir(blazeSpawnsPath).createDirectoryAndParents();
+    CgroupsInfo blazeSpawnsCgroup =
+        new CgroupsInfoV2(CgroupsInfo.Type.BLAZE_SPAWNS, new File(blazeSpawnsPath));
+
+    CgroupsInfo cgroup = blazeSpawnsCgroup.createIndividualSpawnCgroup("spawn_1", 0);
+
+    assertThat(cgroup.exists()).isTrue();
+    assertThat(cgroup.getCgroupDir().getPath()).isEqualTo(blazeSpawnsPath + "/spawn_1.scope");
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.SPAWN);
+    assertThat(cgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    // In reality, cgroups should still create this file automatically, but since we don't have
+    // that in our tests, the memory limits files should not have been created since they aren't
+    // written to.
+    assertThat(new File(blazeSpawnsPath + "/spawn_1.scope/memory.oom.group").exists()).isFalse();
+    assertThat(new File(blazeSpawnsPath + "/spawn_1.scope/memory.max").exists()).isFalse();
+  }
+
+  @Test
+  public void testGetMemoryUsageInKb_v1() throws IOException {
+    String cgroupPath = root + "/dev/cgroup/memory/memdir/action-1";
+    scratch.dir(cgroupPath).createDirectoryAndParents();
+    CgroupsInfo cgroupsInfo =
+        new CgroupsInfoV1(CgroupsInfo.Type.SPAWN, /* cgroupDir= */ new File(cgroupPath));
+
+    // It should return 0 if the /path/to/cgroup/memory.usage_in_bytes does not exist.
+    assertThat(cgroupsInfo.getMemoryUsageInKb()).isEqualTo(0);
+
+    scratch.file(cgroupPath + "/memory.usage_in_bytes", "1024000");
+
+    assertThat(cgroupsInfo.getMemoryUsageInKb()).isEqualTo(1000);
+  }
+
+  @Test
+  public void testGetMemoryUsageInKb_v2() throws IOException {
+    String cgroupPath = root + "/sys/fs/cgroup/memdir/action-1";
+    scratch.dir(cgroupPath).createDirectoryAndParents();
+    CgroupsInfo cgroupsInfo =
+        new CgroupsInfoV2(CgroupsInfo.Type.SPAWN, /* cgroupDir= */ new File(cgroupPath));
+
+    // It should return 0 if the /path/to/cgroup/memory.current does not exist.
+    assertThat(cgroupsInfo.getMemoryUsageInKb()).isEqualTo(0);
+
+    scratch.file(cgroupPath + "/memory.current", "1024000");
+    // Divided by 1024.
+
+    assertThat(cgroupsInfo.getMemoryUsageInKb()).isEqualTo(1000);
+  }
+
+  @Test
+  public void testAddProcess_v1() throws IOException {
+    String cgroupPath = root + "/dev/cgroup/memory/memdir/action-1";
+    scratch.dir(cgroupPath).createDirectoryAndParents();
+    CgroupsInfo cgroupsInfo =
+        new CgroupsInfoV1(CgroupsInfo.Type.SPAWN, /* cgroupDir= */ new File(cgroupPath));
+
+    cgroupsInfo.addProcess(1234);
+
+    assertThat(Files.readLines(new File(cgroupsInfo.getCgroupDir(), "cgroup.procs"), UTF_8))
+        .containsExactly("1234");
+  }
+
+  @Test
+  public void testAddProcess_v2() throws IOException {
+    String cgroupPath = root + "/sys/fs/cgroup/memdir/action-1";
+    scratch.dir(cgroupPath).createDirectoryAndParents();
+    CgroupsInfo cgroupsInfo =
+        new CgroupsInfoV2(CgroupsInfo.Type.SPAWN, /* cgroupDir= */ new File(cgroupPath));
+
+    cgroupsInfo.addProcess(1234);
+
+    assertThat(Files.readLines(new File(cgroupsInfo.getCgroupDir(), "cgroup.procs"), UTF_8))
+        .containsExactly("1234");
+  }
+
+  @Test
+  public void testGetRootCgroup_returnsInvalidCgroup_whenMountNotFound() throws IOException {
+    String pathString = createFakeAbsoluteFile("/proc/self/mounts", "");
+
+    CgroupsInfo cgroup = CgroupsInfo.getRootCgroup(new File(pathString));
+
+    assertThat(cgroup.getClass()).isEqualTo(InvalidCgroupsInfo.class);
+    assertThat(cgroup.getType()).isEqualTo(CgroupsInfo.Type.ROOT);
+  }
+
+  @Test
+  public void testCreateCgroupFromInvalidCgroup_returnsInvalidCgroup() {
+    String errorMessage = "Some error message";
+    CgroupsInfo invalidRootCgroup =
+        new InvalidCgroupsInfo(CgroupsInfo.Type.ROOT, CgroupsInfo.Version.V1, errorMessage);
+    CgroupsInfo invalidBlazeSpawnsCgroup =
+        new InvalidCgroupsInfo(CgroupsInfo.Type.BLAZE_SPAWNS, CgroupsInfo.Version.V2, errorMessage);
+
+    CgroupsInfo createdBlazeSpawnsCgroup = invalidRootCgroup.createBlazeSpawnsCgroup("");
+    CgroupsInfo createdSpawnCgroup =
+        invalidBlazeSpawnsCgroup.createIndividualSpawnCgroup("spawn_1", 1);
+
+    assertThat(createdBlazeSpawnsCgroup.getClass()).isEqualTo(InvalidCgroupsInfo.class);
+    // Should still have the same version as the parent cgroup that attempted to create it.
+    assertThat(createdBlazeSpawnsCgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V1);
+    assertThat(createdBlazeSpawnsCgroup.getType()).isEqualTo(CgroupsInfo.Type.BLAZE_SPAWNS);
+
+    assertThat(createdSpawnCgroup.getClass()).isEqualTo(InvalidCgroupsInfo.class);
+    // Should still have the same version as the parent cgroup that attempted to create it.
+    assertThat(createdSpawnCgroup.getVersion()).isEqualTo(CgroupsInfo.Version.V2);
+    assertThat(createdSpawnCgroup.getType()).isEqualTo(CgroupsInfo.Type.SPAWN);
+  }
+
   private String createFakeAbsoluteFile(String fileName, String... contents) throws IOException {
     return scratch.file(root + fileName, contents).getPathString();
   }