rocker/src/cgroup.rs

use std::fmt::{Display, Formatter};
use nix::sys::signal::kill;
use nix::unistd::Pid;
use nix::sys::signal::Signal;
use std::path::Path;
use crate::{create_dir, get_container_info, ContainerInfo};
use crate::error::Result;

static CGROUP_PATH: &str = "/sys/fs/cgroup/rocker";

pub enum CgroupLevel {
    One,
    Two,
    Three,
}

impl From<&str> for CgroupLevel {
    fn from(value: &str) -> Self {
        match value {
            "1" => CgroupLevel::One,
            "2" => CgroupLevel::Two,
            "3" => CgroupLevel::Three,
            _ => todo!("invalid cgroup level"),
        }
    }
}

impl Display for CgroupLevel {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::One => write!(f, "1"),
            Self::Two => write!(f, "2"),
            Self::Three => write!(f, "3"),
        }
    }
}

impl Default for CgroupLevel {
    fn default() -> Self {
        CgroupLevel::One
    }
}

pub struct Cgroup {
    pub procs: Vec<i32>
}

impl From<&ContainerInfo> for Cgroup {
    fn from(container_inro: &ContainerInfo) -> Self {
        // let container_inro = get_container_info(container_id).unwrap();
        let l = CgroupLevel::from(container_inro.cgroup_level.as_str());
        let cgroup_path = format!("{CGROUP_PATH}_{l}/{}", container_inro.id);
        let mut v = vec![];
        for pid in std::fs::read_to_string(format!("{cgroup_path}/cgroup.procs")).unwrap_or_default().lines() {
            let pid = pid.parse::<i32>().unwrap();
            v.push(pid);
        }
        Self {
            procs: v
        }
    }
}

impl Cgroup {
    pub fn create(container_id: &str, pid: i32, l: CgroupLevel) -> Result<()> {
        let cgroup_path = format!("{CGROUP_PATH}_{l}/{container_id}");
        create_dir(&cgroup_path, false)?;
        // 打开 cgroup.procs, 写入pid
        std::fs::write(format!("{cgroup_path}/cgroup.procs"), pid.to_string())?;
        Ok(())
    }

    pub fn remove(container_id: &str) {
        let container_inro = get_container_info(container_id).unwrap();
        let _self = Self::from(&container_inro);
        for pid in _self.procs {
            let _ = kill(Pid::from_raw(pid), Signal::SIGKILL);
            let pid_path = Path::new("/proc").join(pid.to_string());
            (0..100).any(|i|{
                std::thread::sleep(std::time::Duration::from_millis(i));
                pid_path.exists() == false
            });
        }
        let l = container_inro.cgroup_level;
        let container_group_path = format!("{CGROUP_PATH}_{l}/{container_id}");
        let _ = std::fs::remove_dir(&container_group_path);
        println!("删除cgroup {container_group_path}")
    }
}