diff --git a/Cargo.lock b/Cargo.lock index 1eba973..7ce0d46 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -171,6 +171,12 @@ version = "3.20.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" +[[package]] +name = "bytecount" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "175812e0be2bccb6abe50bb8d566126198344f707e304f45c648fd8f2cc0365e" + [[package]] name = "bytes" version = "1.11.1" @@ -356,6 +362,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", + "tabled", "thiserror", "tokio", ] @@ -421,6 +428,12 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "foldhash" version = "0.1.5" @@ -1106,6 +1119,17 @@ dependencies = [ "winapi", ] +[[package]] +name = "papergrid" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6978128c8b51d8f4080631ceb2302ab51e32cc6e8615f735ee2f83fd269ae3f1" +dependencies = [ + "bytecount", + "fnv", + "unicode-width", +] + [[package]] name = "paste" version = "1.0.15" @@ -1170,6 +1194,28 @@ dependencies = [ "syn", ] +[[package]] +name = "proc-macro-error-attr2" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96de42df36bb9bba5542fe9f1a054b8cc87e172759a1868aa05c1f3acc89dfc5" +dependencies = [ + "proc-macro2", + "quote", +] + +[[package]] +name = "proc-macro-error2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11ec05c52be0a07b08061f7dd003e7d7092e0472bc731b4af7bb1ef876109802" +dependencies = [ + "proc-macro-error-attr2", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -1509,6 +1555,30 @@ dependencies = [ "syn", ] +[[package]] +name = "tabled" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e39a2ee1fbcd360805a771e1b300f78cc88fec7b8d3e2f71cd37bbf23e725c7d" +dependencies = [ + "papergrid", + "tabled_derive", + "testing_table", +] + +[[package]] +name = "tabled_derive" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ea5d1b13ca6cff1f9231ffd62f15eefd72543dab5e468735f1a456728a02846" +dependencies = [ + "heck", + "proc-macro-error2", + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "tempfile" version = "3.27.0" @@ -1522,6 +1592,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "testing_table" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f8daae29995a24f65619e19d8d31dea5b389f3d853d8bf297bbf607cd0014cc" +dependencies = [ + "unicode-width", +] + [[package]] name = "thiserror" version = "2.0.18" @@ -1678,6 +1757,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + [[package]] name = "unicode-xid" version = "0.2.6" diff --git a/Cargo.toml b/Cargo.toml index 9e57756..e76dc4e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ rust-argon2 = "3.0.0" serde = { version = "1.0.219", features = ["derive"] } serde_json = "1.0.140" serde_yaml = "0.9.34" +tabled = "0.20.0" thiserror = "2.0.12" tokio = { version = "1.45.1", features = ["fs", "io-std", "macros", "process", "rt"] } diff --git a/src/bin/dkl.rs b/src/bin/dkl.rs index e9e6de3..e0053b3 100644 --- a/src/bin/dkl.rs +++ b/src/bin/dkl.rs @@ -3,6 +3,7 @@ use eyre::{format_err, Result}; use human_units::Duration; use log::{debug, error}; use std::net::SocketAddr; +use std::path::PathBuf; use tokio::fs; #[derive(Parser)] @@ -81,6 +82,21 @@ enum Command { #[arg(long, default_value = "5s")] timeout: Duration, }, + + Cg { + #[command(subcommand)] + cmd: CgCmd, + }, +} + +#[derive(Subcommand)] +enum CgCmd { + Ls { + #[arg(long)] + root: Option, + #[arg(long, short = 'X')] + exclude: Vec, + }, } #[tokio::main(flavor = "current_thread")] @@ -157,6 +173,10 @@ async fn main() -> Result<()> { .run() .await .map(|_| ())?), + + C::Cg { cmd } => match cmd { + CgCmd::Ls { root, exclude } => Ok(dkl::cgroup::ls(root, &exclude).await?), + }, } } diff --git a/src/cgroup.rs b/src/cgroup.rs new file mode 100644 index 0000000..81232cc --- /dev/null +++ b/src/cgroup.rs @@ -0,0 +1,358 @@ +use std::borrow::Cow; +use std::fmt::Display; +use std::path::{Path as StdPath, PathBuf}; +use std::rc::Rc; +use std::str::FromStr; + +use crate::{fs, human::Human}; + +const CGROUP_ROOT: &str = "/sys/fs/cgroup/"; + +pub async fn ls(parent: Option>, exclude: &[String]) -> fs::Result<()> { + let mut root = PathBuf::from(CGROUP_ROOT); + if let Some(parent) = parent { + root = root.join(parent); + } + + let mut todo = vec![(Cgroup::root(root).await?, vec![], true)]; + + let mut table = tabled::builder::Builder::new(); + table.push_record(["cgroup", "workg set", "anon", "max"]); + + while let Some((cg, p_lasts, last)) = todo.pop() { + let mut name = String::new(); + for last in p_lasts.iter().skip(1) { + name.push_str(if *last { " " } else { "| " }); + } + if !p_lasts.is_empty() { + name.push_str(if last { "`- " } else { "|- " }); + } + name.push_str(&cg.name()); + + table.push_record([ + name, + cg.memory.working_set().human(), + cg.memory.stat.anon.human(), + cg.memory.max.human(), + ]); + + let mut p_lasts = p_lasts.clone(); + p_lasts.push(last); + + let mut children = cg.read_children().await?; + children.sort(); + todo.extend( + (children.into_iter().rev()) + .filter(|c| !exclude.iter().any(|x| x == &c.path.full_name())) + .enumerate() + .map(|(i, child)| (child, p_lasts.clone(), i == 0)), + ); + } + + use tabled::settings::{ + object::{Column, Row}, + Alignment, Modify, + }; + let mut table = table.build(); + table.with(tabled::settings::Style::psql()); + table.with(Alignment::right()); + table.with(Modify::list(Column::from(0), Alignment::left())); + table.with(Modify::list(Row::from(0), Alignment::left())); + + println!("{}", table); + + Ok(()) +} + +pub struct Cgroup { + path: Rc, + children: Vec, + memory: Memory, +} + +impl Cgroup { + pub async fn root(path: impl AsRef) -> fs::Result { + let path = path.as_ref(); + Self::read(Path::root(path), path).await + } + + async fn read(cg_path: Rc, path: impl AsRef) -> fs::Result { + let path = path.as_ref(); + + use fs::Error as E; + + let mut rd = fs::read_dir(path).await?; + + let mut cg = Self { + path: cg_path, + children: Vec::new(), + memory: Memory::default(), + }; + + while let Some(entry) = (rd.next_entry().await).map_err(|e| E::ReadDir(path.into(), e))? { + let path = entry.path(); + + let Some(file_name) = path.file_name() else { + continue; + }; + + if (entry.file_type().await) + .map_err(|e| E::Stat(path.clone(), e))? + .is_dir() + { + cg.children.push(file_name.into()); + continue; + } + + let file_name = file_name.as_encoded_bytes(); + let Some(idx) = file_name.iter().position(|b| *b == b'.') else { + continue; + }; + + let (controller, param) = file_name.split_at(idx); + let param = ¶m[1..]; + + match controller { + b"memory" => match param { + b"current" => cg.memory.current = read_parse(path).await?, + b"low" => cg.memory.low = read_parse(path).await?, + b"high" => cg.memory.high = read_parse(path).await?, + b"min" => cg.memory.min = read_parse(path).await?, + b"max" => cg.memory.max = read_parse(path).await?, + b"stat" => cg.memory.stat.read_from(path).await?, + _ => {} + }, + _ => {} + } + } + + Ok(cg) + } + + async fn read_children(&self) -> fs::Result> { + let mut r = Vec::with_capacity(self.children.len()); + + let mut dir = PathBuf::from(self.path.as_ref()); + + for child_name in &self.children { + dir.push(child_name); + let child_path = Path::Child(self.path.clone(), child_name.into()); + r.push(Self::read(child_path.into(), &dir).await?); + dir.pop(); + } + + Ok(r) + } + + pub fn name(&self) -> Cow<'_, str> { + self.path.name() + } + + pub fn full_name(&self) -> String { + self.path.full_name() + } + + pub fn children(&self) -> impl Iterator { + self.children.iter().map(|n| n.as_path()) + } + + pub async fn read_child(&self, name: impl AsRef) -> fs::Result { + let name = name.as_ref(); + let mut dir = PathBuf::from(self.path.as_ref()); + dir.push(name); + Self::read(self.path.child(name), &dir).await + } +} + +impl PartialEq for Cgroup { + fn eq(&self, o: &Self) -> bool { + &self.path == &o.path + } +} +impl Eq for Cgroup {} + +impl Ord for Cgroup { + fn cmp(&self, o: &Self) -> std::cmp::Ordering { + self.path.cmp(&o.path) + } +} +impl PartialOrd for Cgroup { + fn partial_cmp(&self, o: &Self) -> Option { + Some(self.cmp(o)) + } +} + +#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)] +enum Path { + Root(PathBuf), + Child(Rc, PathBuf), +} + +impl Path { + fn name(&self) -> Cow<'_, str> { + match self { + Self::Root(_) => "/".into(), + Self::Child(_, n) => n.to_string_lossy(), + } + } + + fn full_name(&self) -> String { + use Path::*; + match self { + Root(_) => "/".into(), + Child(parent, _) => match parent.as_ref() { + Root(_) => self.name().into(), + Child(_, _) => format!("{}/{}", parent.full_name(), self.name()), + }, + } + } + + fn depth(&self) -> usize { + use Path::*; + match self { + Root(_) => 0, + Child(p, _) => 1 + p.depth(), + } + } + + fn root(dir: impl Into) -> Rc { + Rc::new(Self::Root(dir.into())) + } + + fn child(self: &Rc, name: impl Into) -> Rc { + Rc::new(Self::Child(self.clone(), name.into())) + } +} + +impl From<&Path> for PathBuf { + fn from(mut p: &Path) -> Self { + let mut stack = Vec::with_capacity(p.depth() + 1); + loop { + match p { + Path::Root(root_path) => { + stack.push(root_path); + break; + } + Path::Child(parent, n) => { + stack.push(n); + p = parent; + } + } + } + + let len = stack.iter().map(|p| p.as_os_str().len() + 1).sum::() - 1; + + let mut buf = PathBuf::with_capacity(len); + buf.extend(stack.into_iter().rev()); + buf + } +} + +#[test] +fn test_path_to_pathbuf() { + let root = Path::root("/a/b"); + let c1 = root.child("c1"); + let c1_1 = c1.child("c1-1"); + + assert_eq!(PathBuf::from("/a/b/c1"), PathBuf::from(c1.as_ref())); + assert_eq!(PathBuf::from("/a/b/c1/c1-1"), PathBuf::from(c1_1.as_ref())); +} + +#[derive(Default)] +struct Memory { + current: Option, + low: Option, + high: Option, + min: Option, + max: Option, + stat: MemoryStat, +} + +impl Memory { + /// working set as defined by cAdvisor + /// (https://github.com/google/cadvisor/blob/e1ccfa9b4cf2e17d74e0f5526b6487b74b704503/container/libcontainer/handler.go#L853-L862) + fn working_set(&self) -> Option { + let cur = self.current?; + let inactive = self.stat.inactive_file?; + (inactive <= cur).then(|| cur - inactive) + } +} + +#[derive(Default)] +struct MemoryStat { + anon: Option, + file: Option, + kernel: Option, + kernel_stack: Option, + pagetables: Option, + shmem: Option, + inactive_file: Option, +} + +enum Max { + Num(u64), + Max, +} + +impl Display for Max { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Num(n) => write!(f, "{n}"), + Self::Max => f.write_str("max"), + } + } +} + +impl Human for Max { + fn human(&self) -> String { + match self { + Self::Num(n) => n.human(), + Self::Max => "+∞".into(), + } + } +} + +impl FromStr for Max { + type Err = std::num::ParseIntError; + fn from_str(s: &str) -> std::result::Result::Err> { + Ok(match s { + "max" => Self::Max, + s => Self::Num(s.parse()?), + }) + } +} + +async fn read_parse(path: impl AsRef) -> fs::Result> +where + T::Err: Display, +{ + let path = path.as_ref(); + + (fs::read_to_string(path).await?) + .trim_ascii() + .parse() + .map_err(|e| fs::Error::Other(path.into(), format!("parse failed: {e}"))) + .map(|v| Some(v)) +} + +impl MemoryStat { + async fn read_from(&mut self, path: impl AsRef) -> fs::Result<()> { + for line in (fs::read_to_string(path).await?).lines() { + let Some((key, value)) = line.split_once(' ') else { + continue; + }; + let value = value.parse::().ok(); + match key { + "anon" => self.anon = value, + "file" => self.file = value, + "kernel" => self.kernel = value, + "kernel_stack" => self.kernel_stack = value, + "pagetables" => self.pagetables = value, + "shmem" => self.shmem = value, + "inactive_file" => self.inactive_file = value, + _ => {} + } + } + Ok(()) + } +} diff --git a/src/fs.rs b/src/fs.rs index 2fbb323..7d449fb 100644 --- a/src/fs.rs +++ b/src/fs.rs @@ -1,9 +1,38 @@ -use eyre::Result; use std::fs::Metadata; -use std::path::PathBuf; -use tokio::fs::read_dir; +use std::path::{Path, PathBuf}; +use tokio::fs; use tokio::sync::mpsc; +pub type Result = std::result::Result; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("{0}: read_dir: {1}")] + ReadDir(PathBuf, std::io::Error), + #[error("{0}: read: {1}")] + Read(PathBuf, std::io::Error), + #[error("{0}: stat: {1}")] + Stat(PathBuf, std::io::Error), + #[error("{0}: write: {1}")] + Write(PathBuf, std::io::Error), + #[error("{0}: {1}")] + Other(PathBuf, String), +} + +macro_rules! wrap_path { + ($fn:ident $( ( $( $pname:ident : $ptype:ty ),* ) )? -> $result:ty, $err:ident) => { + pub async fn $fn(path: impl AsRef$($(, $pname: $ptype)*)?) -> Result<$result> { + let path = path.as_ref(); + fs::$fn(path $($(, $pname)*)?).await.map_err(|e| Error::$err(path.into(), e)) + } + }; +} + +wrap_path!(read_dir -> fs::ReadDir, ReadDir); +wrap_path!(read -> Vec, Read); +wrap_path!(read_to_string -> String, Read); +wrap_path!(write(content: &[u8]) -> (), Write); + pub fn spawn_walk_dir( dir: impl Into + Send + 'static, ) -> mpsc::Receiver> { @@ -24,7 +53,7 @@ pub async fn walk_dir(dir: impl Into, tx: mpsc::Sender v, Err(e) => { - if tx.send(Err(e.into())).await.is_err() { + if tx.send(Err(Error::ReadDir(dir.clone(), e))).await.is_err() { return; } todo.pop_front(); // skip dir on error diff --git a/src/human.rs b/src/human.rs new file mode 100644 index 0000000..e1f3ef2 --- /dev/null +++ b/src/human.rs @@ -0,0 +1,35 @@ +use human_units::FormatSize; +use std::fmt::{Display, Formatter, Result}; + +pub trait Human { + fn human(&self) -> String; +} + +pub struct Quantity(u64); + +impl Display for Quantity { + fn fmt(&self, f: &mut Formatter<'_>) -> Result { + self.0.format_size().fmt(f) + } +} + +impl Human for Quantity { + fn human(&self) -> String { + self.to_string() + } +} + +impl Human for u64 { + fn human(&self) -> String { + self.format_size().to_string() + } +} + +impl Human for Option { + fn human(&self) -> String { + match self { + Some(h) => h.human(), + None => "◌".to_string(), + } + } +} diff --git a/src/lib.rs b/src/lib.rs index 7d1bd9e..35cb920 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,7 @@ pub mod apply; +pub mod rc; +pub mod cgroup; +pub mod human; pub mod bootstrap; pub mod dls; pub mod dynlay; diff --git a/src/rc.rs b/src/rc.rs new file mode 100644 index 0000000..e69de29