From 7bdc6a24f8109c5e69c9754e6fbf8d625124dbd6 Mon Sep 17 00:00:00 2001
From: James Sturtevant <jstur@microsoft.com>
Date: Thu, 10 Aug 2023 23:53:57 +0000
Subject: [PATCH] Handle cgroup v1 or v2 to enable resource updates

Signed-off-by: James Sturtevant <jstur@microsoft.com>
---
 .github/workflows/ci.yml  |  2 +-
 crates/shim/src/cgroup.rs | 50 +++++++++++++++++++++------------------
 2 files changed, 28 insertions(+), 24 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 9375e818..64befe9c 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -136,7 +136,7 @@ jobs:
           GOPROXY: direct
           TEST_RUNTIME: "io.containerd.runc.v2-rs"
           TESTFLAGS_PARALLEL: 1
-          EXTRA_TESTFLAGS: "-no-criu -test.skip='(TestContainerPTY|TestContainerExecLargeOutputWithTTY|TestTaskUpdate|TestTaskResize)'"
+          EXTRA_TESTFLAGS: "-no-criu -test.skip='(TestContainerPTY|TestContainerExecLargeOutputWithTTY)'"
           TESTFLAGS_RACE: "-race"
         run: sudo -E PATH=$PATH make integration
         working-directory: src/github.com/containerd/containerd
diff --git a/crates/shim/src/cgroup.rs b/crates/shim/src/cgroup.rs
index b986bd7c..611e2f13 100644
--- a/crates/shim/src/cgroup.rs
+++ b/crates/shim/src/cgroup.rs
@@ -97,26 +97,7 @@ fn write_process_oom_score(pid: u32, score: i64) -> Result<()> {
 pub fn collect_metrics(pid: u32) -> Result<Metrics> {
     let mut metrics = Metrics::new();
 
-    let hierarchies = hierarchies::auto();
-    let cgroup = if hierarchies.v2() {
-        let path = format!("/proc/{}/cgroup", pid);
-        let content = fs::read_to_string(path).map_err(io_error!(e, "read cgroup"))?;
-        let content = content.strip_suffix('\n').unwrap_or_default();
-
-        let parts: Vec<&str> = content.split("::").collect();
-        let path_parts: Vec<&str> = parts[1].split('/').collect();
-        let namespace = path_parts[1];
-        let cgroup_name = path_parts[2];
-        Cgroup::load(
-            hierarchies,
-            format!("/sys/fs/cgroup/{namespace}/{cgroup_name}").as_str(),
-        )
-    } else {
-        // get container main process cgroup
-        let path = get_cgroups_relative_paths_by_pid(pid)
-            .map_err(other_error!(e, "get process cgroup"))?;
-        Cgroup::load_with_relative_paths(hierarchies::auto(), Path::new("."), path)
-    };
+    let cgroup = get_cgroup(pid)?;
 
     // to make it easy, fill the necessary metrics only.
     for sub_system in Cgroup::subsystems(&cgroup) {
@@ -193,12 +174,35 @@ pub fn collect_metrics(pid: u32) -> Result<Metrics> {
     Ok(metrics)
 }
 
+// get_cgroup will return either cgroup v1 or v2 depending on system configuration
+fn get_cgroup(pid: u32) -> Result<Cgroup> {
+    let hierarchies = hierarchies::auto();
+    let cgroup = if hierarchies.v2() {
+        let path = format!("/proc/{}/cgroup", pid);
+        let content = fs::read_to_string(path).map_err(io_error!(e, "read cgroup"))?;
+        let content = content.strip_suffix('\n').unwrap_or_default();
+
+        let parts: Vec<&str> = content.split("::").collect();
+        let path_parts: Vec<&str> = parts[1].split('/').collect();
+        let namespace = path_parts[1];
+        let cgroup_name = path_parts[2];
+        Cgroup::load(
+            hierarchies,
+            format!("/sys/fs/cgroup/{namespace}/{cgroup_name}").as_str(),
+        )
+    } else {
+        // get container main process cgroup
+        let path = get_cgroups_relative_paths_by_pid(pid)
+            .map_err(other_error!(e, "get process cgroup"))?;
+        Cgroup::load_with_relative_paths(hierarchies::auto(), Path::new("."), path)
+    };
+    Ok(cgroup)
+}
+
 /// Update process cgroup limits
 pub fn update_resources(pid: u32, resources: &LinuxResources) -> Result<()> {
     // get container main process cgroup
-    let path =
-        get_cgroups_relative_paths_by_pid(pid).map_err(other_error!(e, "get process cgroup"))?;
-    let cgroup = Cgroup::load_with_relative_paths(hierarchies::auto(), Path::new("."), path);
+    let cgroup = get_cgroup(pid)?;
 
     for sub_system in Cgroup::subsystems(&cgroup) {
         match sub_system {