From 6639b2b028b93b1bd8a776f2a082c93fa0f19acf Mon Sep 17 00:00:00 2001
From: David Porter <porterdavid@google.com>
Date: Fri, 15 Jul 2022 18:29:49 -0700
Subject: [PATCH] test: Add ability to switch between  cgroupv1 or cgroupv2

For k8s e2e node tests, we want to be able to test both cgroupv1 and
cgroupv2 variants. Currently, the latest version of COS (M97) is running
as cgroupv2 and we would like the add the ability to switch to cgroupv1
for running tests on cgroupv1.

To do this, add a new flag, "CGROUP_MODE" which can be set as v1 or v2.
If the current cgroup configuration does not match the desired config,
the cgroup config will be set via `cgroup_helper` COS utility, and the
node will be rebooted.

Signed-off-by: David Porter <porterdavid@google.com>
---
 test/e2e_node/gci-init.sh | 57 +++++++++++++++++++++++++--------------
 1 file changed, 37 insertions(+), 20 deletions(-)

diff --git a/test/e2e_node/gci-init.sh b/test/e2e_node/gci-init.sh
index 83da28194..47ac30c33 100755
--- a/test/e2e_node/gci-init.sh
+++ b/test/e2e_node/gci-init.sh
@@ -16,7 +16,6 @@
 
 # This script is used to do extra initialization on GCI.
 
-# set up cgroupv2 based on flag CONTAINERD_CGROUPV2 in containerd-env
 CONTAINERD_HOME="/home/containerd"
 CONTAINERD_ENV_METADATA="containerd-env"
 
@@ -24,28 +23,46 @@ if [ -f "${CONTAINERD_HOME}/${CONTAINERD_ENV_METADATA}" ]; then
   source "${CONTAINERD_HOME}/${CONTAINERD_ENV_METADATA}"
 fi
 
-if [ "${CONTAINERD_CGROUPV2:-"false"}"  == "true" ]; then
-  # check cos image
-  if [ -r /etc/os-release ]; then
-    OS_ID="$(. /etc/os-release && echo "$ID")"
+# CONTAINERD_COS_CGROUP_MODE can be specified as "v1" or "v2". If specified,
+# cgroup configuration will be switched as appropriate.
+function configure_cgroup_mode() {
+  if [[ -z "${CONTAINERD_COS_CGROUP_MODE}" ]]; then
+    return
   fi
-  if [ "${OS_ID}" = "cos" ]; then
-    if ! grep -q 'systemd.unified_cgroup_hierarchy=true' /proc/cmdline && [ "$(stat -fc %T /sys/fs/cgroup/)" != "cgroup2fs" ]; then
-      echo "Setting up cgroupv2"
 
-      mount_path="/tmp/esp"
-      mkdir -p "${mount_path}"
-      esp_partition="/dev/sda12"
-      mount "${esp_partition}" "${mount_path}"
-      sed -i 's/systemd.unified_cgroup_hierarchy=false/systemd.unified_cgroup_hierarchy=true/g' "${mount_path}/efi/boot/grub.cfg"
-      umount "${mount_path}"
-      rmdir "${mount_path}"
-
-      echo "Reconfigured grub; rebooting..."
-      reboot
-    fi
+  if [[ ! -r /etc/os-release ]]; then
+    echo "Skipped configuring cgroup mode to ${CONTAINERD_COS_CGROUP_MODE} because /etc/os-release was not readable"
+    return
   fi
-fi
+
+  OS_ID="$(cat /etc/os-release | grep '^ID=' | sed -e 's/ID=//')"
+  if [[ "${OS_ID}" != "cos" ]]; then
+    echo "Skipped configuring cgroup mode to ${CONTAINERD_COS_CGROUP_MODE} because OS is not COS"
+    return
+  fi
+
+  # cgroup_helper was introduced in COS M97, see if it's available first...
+  if ! command -v cgroup_helper > /dev/null 2>&1; then
+    echo "Skipped configuring cgroup mode to ${CONTAINERD_COS_CGROUP_MODE} because cgroup_helper tool is not available (only introduced in COS M97)"
+    return
+  fi
+
+  # if cgroup mode requested was v1 but it's currently set as unified (v2),
+  # switch to hybrid (v1) and reboot
+  if [[ "${CONTAINERD_COS_CGROUP_MODE:-}" == "v1" ]] && cgroup_helper show | grep -q 'unified'; then
+    cgroup_helper set hybrid
+    echo "set cgroup config to hybrid, now rebooting..."
+    reboot
+  # if cgroup mode requested was v2 but it's currently set as hybrid (v1),
+  # switch to unified (v2) and reboot
+  elif [[ "${CONTAINERD_COS_CGROUP_MODE:-}" == "v2" ]] && cgroup_helper show | grep -q 'hybrid'; then
+    cgroup_helper set unified
+    echo "set cgroup config to unified, now rebooting..."
+    reboot
+  fi
+}
+
+configure_cgroup_mode
 
 mount /tmp /tmp -o remount,exec,suid
 #TODO(random-liu): Stop docker and remove this docker thing.