diff --git a/Makefile b/Makefile index 4d2b0c0..5dcad25 100644 --- a/Makefile +++ b/Makefile @@ -28,16 +28,19 @@ install: cp images/* $(WORKSPACE)/images/ chown -R rocker:rocker $(WORKSPACE) + @# 设置cgroup.subtree_control的 + mkdir /sys/fs/cgroup/rocker + echo "+cpuset +cpu +io +memory +hugetlb +pids +rdma +misc" > /sys/fs/cgroup/rocker/cgroup.subtree_control + + clean: -./target/debug/rocker --rm all -rocker --rm all -rm -rf $(WORKSPACE)/* -rm -rf /usr/bin/rocker - -rmdir /sys/fs/cgroup/rocker_1/* - -rmdir /sys/fs/cgroup/rocker_2/* - -rmdir /sys/fs/cgroup/rocker_3/* - -rmdir /sys/fs/cgroup/rocker_* + -rmdir /sys/fs/cgroup/rocker/* + -rmdir /sys/fs/cgroup/rocker -@echo "Cleaning up" -@sudo userdel -r $(USER_NAME) diff --git a/src/cgroup.rs b/src/cgroup.rs deleted file mode 100644 index 412473c..0000000 --- a/src/cgroup.rs +++ /dev/null @@ -1,110 +0,0 @@ -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 { - 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 -} - -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::().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}") - } -} - - - diff --git a/src/cgroups.rs b/src/cgroups.rs new file mode 100644 index 0000000..69c7fe0 --- /dev/null +++ b/src/cgroups.rs @@ -0,0 +1,150 @@ +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, + 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::().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::() + .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::()) { + 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}") + } +} + + + diff --git a/src/main.rs b/src/main.rs index 3169a66..fb6601e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,8 +19,8 @@ use error::{Result, RockerError}; mod error; use network::{create_network, remove_network}; mod network; -use cgroup::{Cgroup, CgroupLevel}; -mod cgroup; +use cgroups::{Cgroup, CgroupLevel}; +mod cgroups; static WORKSPACE: &str = "/home/rocker"; static USER_NAME: &str = "rocker"; @@ -56,10 +56,10 @@ struct RockerArgs { #[arg(long)] image: Option, #[arg(long)] - // --volume "/tmp/test1:tmp/test1,/tmp/test2:tmp/test2" + // --volume "/tmp/test1:tmp/test1 /tmp/test2:tmp/test2" volume: Option, #[arg(long)] - // --env "a=1,b=2,c=3" + // --env "a=1 b=2 c=3" env: Option, // --run /bin/bash --exec container_id @@ -85,11 +85,11 @@ struct RockerArgs { #[arg(long)] root: bool, - // rm "container_id_1, container_id_2, container_id_3" + // rm "container_id_1 container_id_2 container_id_3" #[arg(long)] rm: Option, - // stop "container_id_1, container_id_2, container_id_3" + // stop "container_id_1 container_id_2 container_id_3" #[arg(long)] stop: Option, @@ -97,9 +97,9 @@ struct RockerArgs { #[arg(long)] restart: Option, + // // cgroup "cpu=20 memroy=20" cpu单位% memory单位m #[arg(long)] cgroup: Option, - } @@ -201,7 +201,7 @@ fn init_container_dev>(container_merged_path: P) -> Result<()> { } fn init_container_custom_volume>(container_merged_path: P, custom_volume_s: &String) -> Result<()> { - for custom_volume in custom_volume_s.split(",") { + for custom_volume in custom_volume_s.split(" ") { let custom_volume_v = custom_volume.split(":").collect::>(); if custom_volume_v.len() < 2 { return Err(RockerError::OtherError(format!("volume 参数格式不正确: {custom_volume}"))); @@ -236,7 +236,7 @@ fn get_env_vec(env: &String) -> Result> { let env_text = fs::read_to_string(env)?; env_text.lines().map(String::from).collect::>() } else { - env.split(",").map(String::from).collect::>() + env.split(" ").map(String::from).collect::>() }; Ok(env_vec) } @@ -373,7 +373,7 @@ fn start(container_info: &ContainerInfo, cb: CloneCb, clong_flags: CloneFlags, c Ok(child_pid) => { // exec之前的步骤 create_network(&container_info.id, child_pid.as_raw()); - if let Err(e) = Cgroup::create(&container_info.id, child_pid.as_raw(), CgroupLevel::from(container_info.cgroup_level.as_str())) { + if let Err(e) = Cgroup::create(&container_info.id, child_pid.as_raw(), &container_info.cgroup) { println!("cgroup create error: {}", e) } @@ -444,7 +444,8 @@ fn run_container(container_info: &ContainerInfo, is_exec_cmd: Option<&String>) - let _cb = if let Some(exec_cmd) = is_exec_cmd { let _cb = || { - init_exec_ns(container_info.procs[0]).unwrap(); + let container_cgroup = Cgroup::from(container_info.id.as_str()); + init_exec_ns(container_cgroup.procs[0]).unwrap(); clear_env(); if container_info.root == false { init_container_user(rocker_uid, rocker_gid).unwrap(); @@ -546,14 +547,12 @@ struct ContainerInfo { run: String, // /bin/bash image: String, // busybox - volume: String, // /root/tmp:/root/tmp,/root/tmp1:/root/tmp1 + volume: String, // /root/tmp:/root/tmp /root/tmp1:/root/tmp1 ... env: String, // a=1,b=2,c=3 或者 env文件路径 log: bool, wait: bool, root: bool, - - cgroup_level: String, // v1/v2/v3... - procs: Vec // 动态从控制组中获取的 + cgroup: String, // "cpu=20 memroy=20 ..." cpu单位%, memory单位m } impl ContainerInfo { @@ -582,10 +581,7 @@ fn get_container_info(container_id: &str) -> Result { let mut container_info: ContainerInfo = toml::from_str(&info_str)?; // todo 判断当前进程组 是否还有进程 - for pid in Cgroup::from(&container_info).procs { - container_info.procs.push(pid) - } - if container_info.procs.len() > 0 { + if Cgroup::from(container_id).populated > 0 { container_info.status = ContainerStatus::RUNNING; } else { container_info.status = ContainerStatus::STOP; @@ -701,6 +697,7 @@ fn main() -> Result<()>{ let args = RockerArgs::parse(); if args.image.is_some() || args.restart.is_some() { + // --run / --restart let container_info = match (&args.run, &args.image, &args.restart) { (Some(cmd), Some(image), None) => { let container_id = args.id.unwrap_or(uuid::Uuid::new_v4().to_string()[0..8].to_string()); @@ -717,9 +714,8 @@ fn main() -> Result<()>{ status: ContainerStatus::READY, log: args.log, wait: args.wait, - cgroup_level: args.cgroup.clone().unwrap_or(CgroupLevel::default().to_string()), root: args.root, - procs: vec![] + cgroup: args.cgroup.clone().unwrap_or_default(), } } (None, None, Some(_container_id)) => { diff --git a/test/test.py b/test/test.py index be6b682..6c54d3e 100755 --- a/test/test.py +++ b/test/test.py @@ -1,7 +1,3 @@ -# import urllib.request - -# from datetime import datetime -# import time import mmap import os @@ -13,8 +9,11 @@ print(cnt, "hello", time.time()) while True: cnt += 1 - print(cnt, "hello", time.time()) - time.sleep(0.01) + # if cnt % 1000_0000 == 0: + # print(cnt, "hello", time.time()) + + # print(cnt, "hello", time.time()) + # time.sleep(0.01)