liboj-cgroups 0.1.0

use std::fs;
use std::io;
use std::ops::Deref;
use std::path::Path;
use std::str::FromStr;
use std::time::Duration;

use lazy_static::lazy_static;
use nix::unistd;

use super::Controller;
use crate::{
    attr_file::{ReadAttr, ResetAttr, StatMap},
    hierarchy::HierarchyNode,
};

lazy_static! {
    static ref USER_HZ: usize = {
        match unistd::sysconf(unistd::SysconfVar::CLK_TCK) {
            Ok(Some(user_hz)) => user_hz as usize,
            Ok(None) => panic!("Failed to get USER_HZ"),
            Err(e) => panic!("Failed to get USER_HZ. {}", e),
        }
    };
}

pub trait CpuAcctController: Controller {
    fn reset(&self) -> Box<dyn '_ + ResetAttr>;
    fn stat(&self) -> Box<dyn '_ + ReadAttr<Stat>>;
    fn usage(&self) -> Box<dyn '_ + ReadAttr<Duration>>;
    fn usage_all(&self) -> Box<dyn '_ + ReadAttr<UsageAll>>;
    fn usage_user(&self) -> Box<dyn '_ + ReadAttr<Duration>>;
    fn usage_sys(&self) -> Box<dyn '_ + ReadAttr<Duration>>;
    fn usage_percpu(&self) -> Box<dyn '_ + ReadAttr<Vec<Duration>>>;
    fn usage_percpu_user(&self) -> Box<dyn '_ + ReadAttr<Vec<Duration>>>;
    fn usage_percpu_sys(&self) -> Box<dyn '_ + ReadAttr<Vec<Duration>>>;
}

impl CpuAcctController for HierarchyNode {
    fn reset(&self) -> Box<dyn '_ + ResetAttr> {
        let file = self.as_path().join("cpuacct.usage");
        Box::new(UsageFile(file))
    }

    fn stat(&self) -> Box<dyn '_ + ReadAttr<Stat>> {
        let file = self.as_path().join("cpuacct.stat");
        Box::new(StatFile(file))
    }

    fn usage(&self) -> Box<dyn '_ + ReadAttr<Duration>> {
        let file = self.as_path().join("cpuacct.usage");
        Box::new(UsageFile(file))
    }

    fn usage_all(&self) -> Box<dyn '_ + ReadAttr<UsageAll>> {
        let file = self.as_path().join("cpuacct.usage_all");
        Box::new(UsageAllFile(file))
    }

    fn usage_user(&self) -> Box<dyn '_ + ReadAttr<Duration>> {
        let file = self.as_path().join("cpuacct.usage_user");
        Box::new(UsageFile(file))
    }

    fn usage_sys(&self) -> Box<dyn '_ + ReadAttr<Duration>> {
        let file = self.as_path().join("cpuacct.usage_sys");
        Box::new(UsageFile(file))
    }

    fn usage_percpu(&self) -> Box<dyn '_ + ReadAttr<Vec<Duration>>> {
        let file = self.as_path().join("cpuacct.usage_percpu");
        Box::new(PerCpuFile(file))
    }

    fn usage_percpu_user(&self) -> Box<dyn '_ + ReadAttr<Vec<Duration>>> {
        let file = self.as_path().join("cpuacct.usage_percpu_user");
        Box::new(PerCpuFile(file))
    }

    fn usage_percpu_sys(&self) -> Box<dyn '_ + ReadAttr<Vec<Duration>>> {
        let file = self.as_path().join("cpuacct.usage_percpu_sys");
        Box::new(PerCpuFile(file))
    }
}

struct StatFile<P: AsRef<Path>>(P);

impl<P: AsRef<Path>> ReadAttr<Stat> for StatFile<P> {
    fn read(&self) -> io::Result<Stat> {
        let s = fs::read_to_string(&self.0)?;
        let stat_map = StatMap::from(s.as_str());
        let mut res = Stat {
            user: Duration::from_secs(0),
            system: Duration::from_secs(0),
        };
        match stat_map.get("user").map(|s| s.parse::<f32>()) {
            Some(Ok(usage)) => res.user = Duration::from_secs_f32(usage / *USER_HZ as f32),
            Some(Err(_)) | None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse user cpu usage in cpuacct.stat",
                ));
            }
        }
        match stat_map.get("system").map(|s| s.parse::<f32>()) {
            Some(Ok(usage)) => res.system = Duration::from_secs_f32(usage / *USER_HZ as f32),
            Some(Err(_)) | None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse system cpu usage in cpuacct.stat",
                ))
            }
        }
        Ok(res)
    }
}

pub struct Stat {
    pub user: Duration,
    pub system: Duration,
}

struct UsageFile<P: AsRef<Path>>(P);

impl<P: AsRef<Path>> ReadAttr<Duration> for UsageFile<P> {
    fn read(&self) -> io::Result<Duration> {
        match fs::read_to_string(&self.0)?.trim().parse() {
            Ok(usage) => Ok(Duration::from_nanos(usage)),
            Err(_) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "failed to parse cpu usage in cgroup",
            )),
        }
    }
}

impl<P: AsRef<Path>> ResetAttr for UsageFile<P> {
    fn reset(&self) -> io::Result<()> {
        fs::write(&self.0, b"0")
    }
}

struct UsageAllFile<P: AsRef<Path>>(P);

impl<P: AsRef<Path>> ReadAttr<UsageAll> for UsageAllFile<P> {
    fn read(&self) -> io::Result<UsageAll> {
        let file = self.0.as_ref();
        match fs::read_to_string(file)?.parse() {
            Ok(usages) => Ok(usages),
            Err(_) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("failed to parse cpu usage from {}", file.display()),
            )),
        }
    }
}

pub struct UsageAll {
    usages: Vec<Usage>,
}

impl FromStr for UsageAll {
    type Err = io::Error;

    fn from_str(s: &str) -> io::Result<UsageAll> {
        let mut usages = Vec::new();
        // The first line is "cpu user system"
        for line in s.lines().skip(1) {
            match line.parse() {
                Ok(usage) => usages.push(usage),
                Err(e) => return Err(e),
            }
        }
        Ok(UsageAll { usages })
    }
}

impl Deref for UsageAll {
    type Target = Vec<Usage>;

    fn deref(&self) -> &Vec<Usage> {
        &self.usages
    }
}

pub struct Usage {
    pub cpu: usize,
    pub user: Duration,
    pub system: Duration,
}

impl FromStr for Usage {
    type Err = io::Error;

    fn from_str(s: &str) -> io::Result<Usage> {
        let mut values = s.split_whitespace();
        let cpu: usize = match values.next().map(|value| value.parse()) {
            Some(Ok(value)) => value,
            Some(Err(_)) => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse cpu index when parsing cpu usage.",
                ));
            }
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse cpu usage because cpu index is missing.",
                ));
            }
        };
        let user: Duration = match values.next().map(|value| value.parse()) {
            Some(Ok(value)) => Duration::from_nanos(value),
            Some(Err(_)) => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse user cpu usage when parsing cpu usage.",
                ));
            }
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse cpu usage because user cpu usage is missing.",
                ));
            }
        };
        let system: Duration = match values.next().map(|value| value.parse()) {
            Some(Ok(value)) => Duration::from_nanos(value),
            Some(Err(_)) => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse system cpu usage when parsing cpu usage.",
                ));
            }
            None => {
                return Err(io::Error::new(
                    io::ErrorKind::InvalidData,
                    "failed to parse cpu usage because system cpu usage is missing.",
                ));
            }
        };
        Ok(Usage { cpu, user, system })
    }
}

struct PerCpuFile<P: AsRef<Path>>(P);

impl<P: AsRef<Path>> ReadAttr<Vec<Duration>> for PerCpuFile<P> {
    fn read(&self) -> io::Result<Vec<Duration>> {
        let file = self.0.as_ref();
        let mut usages = Vec::new();
        for usage in fs::read_to_string(&file)?.split_whitespace() {
            match usage.parse() {
                Ok(usage) => usages.push(Duration::from_nanos(usage)),
                Err(_) => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        format!("{} does not contains a valid cpu usage", file.display()),
                    ))
                }
            }
        }
        Ok(usages)
    }
}