wasm: add wasmEdge container process to cgroup

Signed-off-by: Poorunga <[email protected]>

wasm/Cargo.toml

@@ -22,6 +22,7 @@ oci-spec = "0.5.4"
 signal-hook-tokio = { version = "0.3.1", features = ["futures-v0_3"] }
 log = { version = "0.4.17", features = ["std"] }
 time = "0.3.5"
+cgroups-rs = "0.3.2"
 
 wasmedge-sdk = { version = "0.7.1", optional = true }
 

wasm/src/utils.rs

@@ -1,3 +1,4 @@
+use containerd_shim::Error;
 use oci_spec::runtime::Spec;
 
 #[cfg(feature = "wasmedge")]
@@ -75,3 +76,23 @@ pub fn get_memory_limit(spec: &Spec) -> Option<i64> {
         .and_then(|x| x.memory().as_ref())
         .and_then(|x| x.limit())
 }
+
+pub(crate) fn get_rootfs(spec: &Spec) -> containerd_shim::Result<String> {
+    Ok(spec
+        .root()
+        .as_ref()
+        .ok_or(Error::InvalidArgument(
+            "rootfs is not set in runtime spec".to_string(),
+        ))?
+        .path()
+        .display()
+        .to_string())
+}
+
+#[cfg(feature = "wasmedge")]
+pub(crate) fn get_cgroup_path(spec: &Spec) -> Option<String> {
+    spec.linux()
+        .as_ref()
+        .and_then(|linux| linux.cgroups_path().as_ref())
+        .map(|path| path.display().to_string())
+}

wasm/src/wasmedge.rs

@@ -23,6 +23,7 @@ use std::{
     sync::Arc,
 };
 
+use cgroups_rs::{Cgroup, CgroupPid};
 use containerd_shim::{
     api::{CreateTaskRequest, ExecProcessRequest, Status},
     asynchronous::{
@@ -55,7 +56,7 @@ use wasmedge_sdk::{
     params, PluginManager, Vm,
 };
 
-use crate::utils;
+use crate::utils::{get_args, get_cgroup_path, get_envs, get_preopens, get_rootfs};
 
 pub type ExecProcess = ProcessTemplate<WasmEdgeExecLifecycle>;
 pub type InitProcess = ProcessTemplate<WasmEdgeInitLifecycle>;
@@ -115,16 +116,10 @@ impl ContainerFactory<WasmEdgeContainer> for WasmEdgeContainerFactory {
         let mut spec: Spec = read_spec(req.bundle()).await?;
         spec.canonicalize_rootfs(req.bundle())
             .map_err(|e| Error::InvalidArgument(format!("could not canonicalize rootfs: {e}")))?;
-        let rootfs = spec
-            .root()
-            .as_ref()
-            .ok_or(Error::InvalidArgument(
-                "rootfs is not set in runtime spec".to_string(),
-            ))?
-            .path();
-        mkdir(rootfs, 0o711).await?;
+        let rootfs = get_rootfs(&spec)?;
+        mkdir(&rootfs, 0o711).await?;
         for m in req.rootfs() {
-            mount_rootfs(m, rootfs.as_path()).await?
+            mount_rootfs(m, &rootfs).await?
         }
         let stdio = Stdio::new(req.stdin(), req.stdout(), req.stderr(), req.terminal);
         let exit_signal = Arc::new(Default::default());
@@ -160,17 +155,11 @@ impl ProcessLifecycle<InitProcess> for WasmEdgeInitLifecycle {
     async fn start(&self, p: &mut InitProcess) -> containerd_shim::Result<()> {
         let spec = &p.lifecycle.spec;
         let vm = p.lifecycle.prototype_vm.clone();
-        let args = utils::get_args(spec);
-        let envs = utils::get_envs(spec);
-        let rootfs = spec
-            .root()
-            .as_ref()
-            .ok_or(Error::InvalidArgument(
-                "rootfs is not set in runtime spec".to_string(),
-            ))?
-            .path();
-        let mut preopens = vec![format!("/:{}", rootfs.display())];
-        preopens.append(&mut utils::get_preopens(spec));
+        let args = get_args(spec);
+        let envs = get_envs(spec);
+        let rootfs = get_rootfs(spec)?;
+        let mut preopens = vec![format!("/:{}", rootfs)];
+        preopens.append(&mut get_preopens(spec));
 
         debug!(
             "start wasm with args: {:?}, envs: {:?}, preopens: {:?}",
@@ -188,6 +177,18 @@ impl ProcessLifecycle<InitProcess> for WasmEdgeInitLifecycle {
                 p.pid = init_pid;
             }
             ForkResult::Child => {
+                if let Some(cgroup_path) = get_cgroup_path(spec) {
+                    // Add child process to Cgroup
+                    Cgroup::new(
+                        cgroups_rs::hierarchies::auto(),
+                        cgroup_path.trim_start_matches('/'),
+                    )
+                    .and_then(|cgroup| cgroup.add_task(CgroupPid::from(std::process::id() as u64)))
+                    .map_err(other_error!(
+                        e,
+                        format!("failed to add task to cgroup: {}", cgroup_path)
+                    ))?;
+                }
                 match run_wasi_func(vm, args, envs, preopens, p) {
                     Ok(_) => exit(0),
                     // TODO add a pipe? to return detailed error message
@@ -216,7 +217,19 @@ impl ProcessLifecycle<InitProcess> for WasmEdgeInitLifecycle {
         Ok(())
     }
 
-    async fn delete(&self, _p: &mut InitProcess) -> containerd_shim::Result<()> {
+    async fn delete(&self, p: &mut InitProcess) -> containerd_shim::Result<()> {
+        if let Some(cgroup_path) = get_cgroup_path(&p.lifecycle.spec) {
+            // Add child process to Cgroup
+            Cgroup::load(
+                cgroups_rs::hierarchies::auto(),
+                cgroup_path.trim_start_matches('/'),
+            )
+            .delete()
+            .map_err(other_error!(
+                e,
+                format!("failed to delete cgroup: {}", cgroup_path)
+            ))?;
+        }
         Ok(())
     }
 
@@ -263,9 +276,7 @@ impl ProcessLifecycle<ExecProcess> for WasmEdgeExecLifecycle {
     }
 
     async fn delete(&self, _p: &mut ExecProcess) -> containerd_shim::Result<()> {
-        Err(Error::Unimplemented(
-            "exec not supported for wasm containers".to_string(),
-        ))
+        Ok(())
     }
 
     async fn update(

wasm/src/wasmtime.rs

@@ -20,7 +20,7 @@ use tokio::fs::read;
 use wasmtime::{Config, Engine, Extern, Linker, Module, Store, StoreLimits, StoreLimitsBuilder};
 use wasmtime_wasi::{tokio::WasiCtxBuilder, WasiCtx};
 
-use crate::utils;
+use crate::{utils, utils::get_rootfs};
 
 pub type ExecProcess = ProcessTemplate<WasmtimeExecLifecycle>;
 pub type InitProcess = ProcessTemplate<WasmtimeInitLifecycle>;
@@ -59,16 +59,9 @@ impl ContainerFactory<WasmtimeContainer> for WasmtimeContainerFactory {
         let mut spec: Spec = read_spec(req.bundle()).await?;
         spec.canonicalize_rootfs(req.bundle())
             .map_err(|e| Error::InvalidArgument(format!("could not canonicalize rootfs: {e}")))?;
-        let rootfs = spec
-            .root()
-            .as_ref()
-            .ok_or(Error::InvalidArgument(
-                "rootfs is not set in runtime spec".to_string(),
-            ))?
-            .path();
-        mkdir(rootfs, 0o711).await?;
+        let rootfs = get_rootfs(&spec)?;
         for m in req.rootfs() {
-            mount_rootfs(m, rootfs.as_path()).await?
+            mount_rootfs(m, &rootfs).await?
         }
         let stdio = Stdio::new(req.stdin(), req.stdout(), req.stderr(), req.terminal);
         let exit_signal = Arc::new(Default::default());