dkl/src/cgroup.rs

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};

pub const ROOT: &str = "/sys/fs/cgroup";

pub async fn ls(parent: Option<impl AsRef<StdPath>>, exclude: &[String]) -> fs::Result<()> {
    let mut root = PathBuf::from(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::{
        Alignment, Modify,
        object::{Column, Row},
    };
    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<Path>,
    children: Vec<PathBuf>,
    memory: Memory,
}

impl Cgroup {
    pub async fn root(path: impl AsRef<StdPath>) -> fs::Result<Self> {
        let path = path.as_ref();
        Self::read(Path::root(path), path).await
    }

    async fn read(cg_path: Rc<Path>, path: impl AsRef<StdPath>) -> fs::Result<Self> {
        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 = &param[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<Vec<Self>> {
        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<Item = &StdPath> {
        self.children.iter().map(|n| n.as_path())
    }

    pub async fn read_child(&self, name: impl AsRef<StdPath>) -> fs::Result<Self> {
        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<std::cmp::Ordering> {
        Some(self.cmp(o))
    }
}

#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
enum Path {
    Root(PathBuf),
    Child(Rc<Path>, 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<PathBuf>) -> Rc<Self> {
        Rc::new(Self::Root(dir.into()))
    }

    fn child(self: &Rc<Self>, name: impl Into<PathBuf>) -> Rc<Self> {
        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::<usize>() - 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<u64>,
    low: Option<u64>,
    high: Option<Max>,
    min: Option<u64>,
    max: Option<Max>,
    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<u64> {
        let cur = self.current?;
        let inactive = self.stat.inactive_file?;
        (inactive <= cur).then(|| cur - inactive)
    }
}

#[derive(Default)]
struct MemoryStat {
    anon: Option<u64>,
    file: Option<u64>,
    kernel: Option<u64>,
    kernel_stack: Option<u64>,
    pagetables: Option<u64>,
    shmem: Option<u64>,
    inactive_file: Option<u64>,
}

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<Self, <Self as FromStr>::Err> {
        Ok(match s {
            "max" => Self::Max,
            s => Self::Num(s.parse()?),
        })
    }
}

async fn read_parse<T: FromStr>(path: impl AsRef<StdPath>) -> fs::Result<Option<T>>
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<StdPath>) -> 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::<u64>().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(())
    }
}