rocker/src/cgroups.rs

use std::fmt::{Display, Formatter};
use nix::sys::signal::kill;
use nix::unistd::Pid;
use nix::sys::signal::Signal;
use serde::de::value;
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 {
    V1,
    V2,
    V3,
    V4,
    V5,
    V6,
    V7,
    V8,
    V9,
}

impl From<&str> for CgroupLevel {
    fn from(value: &str) -> Self {
        match value {
            "v1" => CgroupLevel::V1,
            "v2" => CgroupLevel::V2,
            "v3" => CgroupLevel::V3,
            "v4" => CgroupLevel::V4,
            "v5" => CgroupLevel::V5,
            "v6" => CgroupLevel::V6,
            "v7" => CgroupLevel::V7,
            "v8" => CgroupLevel::V8,
            "v9" => CgroupLevel::V9,
            _ => todo!("invalid cgroup level"),
        }
    }
}

impl Display for CgroupLevel {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::V1 => write!(f, "v1"),
            Self::V2 => write!(f, "v2"),
            Self::V3 => write!(f, "v3"),
            Self::V4 => write!(f, "v4"),
            Self::V5 => write!(f, "v5"),
            Self::V6 => write!(f, "v6"),
            Self::V7 => write!(f, "v7"),
            Self::V8 => write!(f, "v8"),
            Self::V9 => write!(f, "v9"),
        }
    }
}

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

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

impl From<&str> for Cgroup {
    fn from(container_id: &str) -> Self {
        let cgroup_path = format!("{CGROUP_PATH}/{}", container_id);
        let procs = std::fs::read_to_string(format!("{cgroup_path}/cgroup.procs"))
            .unwrap_or_default()
            .lines()
            .filter_map(|pid|pid.parse::<i32>().ok())
            .collect();

        // 找到 cgroup.events 第一行的值
        let populated = std::fs::read_to_string(format!("{cgroup_path}/cgroup.events"))
            .unwrap_or_default()
            .lines()
            .next()
            .unwrap_or_default()
            .split(" ")
            .nth(1)
            .unwrap_or("0")
            .parse::<i32>()
            .unwrap_or_default();
        
        Self {
            procs,
            populated
        }
    }
}

impl Cgroup {
    pub fn create(container_id: &str, pid: i32, cgroup_s: &str) -> Result<()> {
        let cgroup_path = format!("{CGROUP_PATH}/{container_id}");
        create_dir(&cgroup_path, false)?;

        // 根据 cgroup_s 写入指定参数
        let _ = cgroup_s.split(" ")
            .map(|s|s.split("="))
            .map(|mut x|(x.next(), x.next()))
            .filter_map(|(dev, value)| {
                if let (Some(dev), Ok(value)) = (dev, value.unwrap_or("").parse::<i32>()) {
                    Some((dev, value))
                } else {
                    None
                }
            }).for_each(|(dev, value)| {
                match (dev, value) {
                    ("cpu", 1..=100) => {
                        if let Ok(_) = std::fs::write(format!("{cgroup_path}/cpu.max"), format!("{} 100000", 100000 / 100 * value)) {
                            println!("cpu 限制 {value}% 成功")
                        }
                    }
                    ("memory", 1..=4096) => {
                        
                    }
                    _ => {
                        println!("🔴 没有定义的 cgroup 参数: {dev}={value}")
                    }
                }
            });

        // 加入控制组
        std::fs::write(format!("{cgroup_path}/cgroup.procs"), pid.to_string())?;
        Ok(())
    }

    pub fn remove(container_id: &str) {
        let _self = Self::from(container_id);
        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 cgroup_path = format!("{CGROUP_PATH}/{container_id}");
        let _ = std::fs::remove_dir(&cgroup_path);
        println!("删除cgroup {cgroup_path}")
    }
}