liboj-cgroups 0.1.0

use std::fs;
use std::io;
use std::path::Path;
use std::str::FromStr;

use lazy_static::lazy_static;
use regex::{Regex, RegexSet};

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

pub trait MemoryController: Controller {
    fn usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>>;
    fn limit_in_bytes(&self) -> Box<dyn '_ + AttrFile<Bytes>>;
    fn failcnt(&self) -> Box<dyn '_ + ReadAttr<usize>>;
    fn max_usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>>;
    fn soft_limit_in_bytes(&self) -> Box<dyn '_ + AttrFile<Bytes>>;
    fn stat(&self) -> Box<dyn '_ + ReadAttr<Stat>>;
    fn use_hierarchy(&self) -> Box<dyn '_ + AttrFile<bool>>;
    fn force_empty(&self) -> Box<dyn '_ + ResetAttr>;
    fn swappiness(&self) -> Box<dyn '_ + AttrFile<u8>>;

    fn swap_memory_controller(&self) -> &dyn SwapMemoryController;

    // TODO: Unimplemented method
    // memory.numa_stat
    // memory.oom_control
    // memory.move_charge_at_immigrate
    // memory.pressure_level
}

pub trait SwapMemoryController {
    fn usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>>;
    fn limit_in_bytes(&self) -> Box<dyn '_ + AttrFile<Bytes>>;
    fn failcnt(&self) -> Box<dyn '_ + ReadAttr<usize>>;
    fn max_usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>>;
}

impl MemoryController for HierarchyNode {
    fn usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>> {
        let file = self.as_path().join("memory.usage_in_bytes");
        Box::new(BytesFile(file))
    }

    fn limit_in_bytes(&self) -> Box<dyn '_ + AttrFile<Bytes>> {
        let file = self.as_path().join("memory.limit_in_bytes");
        Box::new(BytesFile(file))
    }

    fn failcnt(&self) -> Box<dyn '_ + ReadAttr<usize>> {
        let file = self.as_path().join("memory.failcnt");
        Box::new(FailCntFile(file))
    }

    fn max_usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>> {
        let file = self.as_path().join("memory.max_usage_in_bytes");
        Box::new(BytesFile(file))
    }

    fn soft_limit_in_bytes(&self) -> Box<dyn '_ + AttrFile<Bytes>> {
        let file = self.as_path().join("memory.soft_limit_in_bytes");
        Box::new(BytesFile(file))
    }

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

    fn use_hierarchy(&self) -> Box<dyn '_ + AttrFile<bool>> {
        let file = self.as_path().join("memory.use_hierarchy");
        Box::new(UseHierarchyFile(file))
    }

    fn force_empty(&self) -> Box<dyn '_ + ResetAttr> {
        let file = self.as_path().join("memory.force_empty");
        Box::new(ForceEmptyFile(file))
    }

    fn swappiness(&self) -> Box<dyn '_ + AttrFile<u8>> {
        let file = self.as_path().join("memory.swappiness");
        Box::new(SwappinessFile(file))
    }

    fn swap_memory_controller(&self) -> &dyn SwapMemoryController {
        self
    }
}

impl SwapMemoryController for HierarchyNode {
    fn usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>> {
        let file = self.as_path().join("memory.memsw.usage_in_bytes");
        Box::new(BytesFile(file))
    }

    fn limit_in_bytes(&self) -> Box<dyn '_ + AttrFile<Bytes>> {
        let file = self.as_path().join("memory.memsw.limit_in_bytes");
        Box::new(BytesFile(file))
    }

    fn failcnt(&self) -> Box<dyn '_ + ReadAttr<usize>> {
        let file = self.as_path().join("memory.memsw.failcnt");
        Box::new(FailCntFile(file))
    }

    fn max_usage_in_bytes(&self) -> Box<dyn '_ + ReadAttr<Bytes>> {
        let file = self.as_path().join("memory.memsw.max_usage_in_bytes");
        Box::new(BytesFile(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 res = stat_attr!(
            stat_map,
            Stat,
            [
                cache,
                rss,
                rss_huge,
                mapped_file,
                pgpgin,
                pgpgout,
                swap,
                dirty,
                writeback,
                inactive_anon,
                active_anon,
                inactive_file,
                active_file,
                unevictable
            ]
        );
        Ok(res)
    }
}

/// Per-memory cgroup local status
pub struct Stat {
    /// of bytes of page cache memory.
    pub cache: usize,
    /// of bytes of anonymous and swap cache memory (includes transparent hugepages).
    pub rss: usize,
    /// of bytes of anonymous transparent hugepages.
    pub rss_huge: usize,
    /// of bytes of mapped file (includes tmpfs/shmem)
    pub mapped_file: usize,
    /// of charging events to the memory cgroup. The charging event happens each time a page is accounted as either mapped anon page(RSS) or cache page(Page Cache) to the cgroup.
    pub pgpgin: usize,
    /// of uncharging events to the memory cgroup. The uncharging event happens each time a page is unaccounted from the cgroup.
    pub pgpgout: usize,
    /// of bytes of swap usage
    pub swap: usize,
    /// of bytes that are waiting to get written back to the disk.
    pub dirty: usize,
    /// of bytes of file/anon cache that are queued for syncing to disk.
    pub writeback: usize,
    /// of bytes of anonymous and swap cache memory on inactive LRU list.
    pub inactive_anon: usize,
    /// of bytes of anonymous and swap cache memory on active LRU list.
    pub active_anon: usize,
    /// of bytes of file-backed memory on inactive LRU list.
    pub inactive_file: usize,
    /// of bytes of file-backed memory on active LRU list.
    pub active_file: usize,
    /// of bytes of memory that cannot be reclaimed (mlocked etc).
    pub unevictable: usize,
}

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

impl<P: AsRef<Path>> ReadAttr<Bytes> for BytesFile<P> {
    fn read(&self) -> io::Result<Bytes> {
        let s = fs::read_to_string(&self.0)?;
        match s.trim().parse() {
            Ok(bytes) => Ok(Bytes::from_bytes(bytes)),
            Err(_) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("failed to parse bytes from {}", self.0.as_ref().display()),
            )),
        }
    }
}

impl<P: AsRef<Path>> WriteAttr<Bytes> for BytesFile<P> {
    fn write(&self, bytes: &Bytes) -> io::Result<()> {
        match bytes.as_bytes() {
            Some(bytes) => fs::write(&self.0, bytes.to_string()),
            None => fs::write(&self.0, b"-1"),
        }
    }
}

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

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

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

#[derive(Clone, Copy, Eq, PartialEq, Debug)]
pub struct Bytes(Option<usize>);

impl Bytes {
    const fn new(bytes: Option<usize>) -> Bytes {
        Bytes(bytes)
    }

    pub const fn infinity() -> Bytes {
        Bytes::new(None)
    }

    pub const fn from_bytes(b: usize) -> Bytes {
        Bytes::new(Some(b))
    }

    pub fn as_bytes(&self) -> Option<usize> {
        self.0
    }
}

/// ```
/// use liboj_cgroups::subsystem::memory::Bytes;
/// const K: usize = 1024;
/// const M: usize = 1048576;
/// const G: usize = 1073741824;
/// assert_eq!(Bytes::from_bytes(K), "1k".parse().unwrap());
/// assert_eq!(Bytes::from_bytes(K), "1K".parse().unwrap());
/// assert_eq!(Bytes::from_bytes(M), "1m".parse().unwrap());
/// assert_eq!(Bytes::from_bytes(M), "1M".parse().unwrap());
/// assert_eq!(Bytes::from_bytes(G), "1g".parse().unwrap());
/// assert_eq!(Bytes::from_bytes(G), "1G".parse().unwrap());
/// assert_eq!(Bytes::infinity(), "-1".parse().unwrap());
/// ```
impl FromStr for Bytes {
    type Err = io::Error;

    #[allow(clippy::trivial_regex)]
    fn from_str(s: &str) -> io::Result<Bytes> {
        lazy_static! {
            static ref BYTES: Regex = Regex::new(r"^(?P<bytes>\d+)[kKmMgG]?$").unwrap();
            static ref BYTES_SET: RegexSet = RegexSet::new(&[
                r"^(\d+)$",
                r"^(\d+)[kK]$",
                r"^(\d+)[mM]$",
                r"^(\d+)[gG]$",
                r"^-1$",
            ])
            .unwrap();
        }
        match BYTES_SET.matches(s).iter().next() {
            Some(i @ 0..=3) => {
                let bytes: usize = BYTES.captures(s).unwrap()["bytes"].parse().unwrap();
                Ok(Bytes::from_bytes(bytes * 1024usize.pow(i as u32)))
            }
            Some(4) => Ok(Bytes::infinity()),
            Some(_) => unreachable!("regex set out of range"),
            None => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("failed to parse bytes \"{}\"", s),
            )),
        }
    }
}

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

impl<P: AsRef<Path>> ReadAttr<bool> for UseHierarchyFile<P> {
    fn read(&self) -> io::Result<bool> {
        let file = self.0.as_ref();
        let s = fs::read_to_string(&file)?;
        match s.trim() {
            "0" => Ok(false),
            "1" => Ok(true),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("failed to parse bool in {}", file.display()),
            )),
        }
    }
}

impl<P: AsRef<Path>> WriteAttr<bool> for UseHierarchyFile<P> {
    fn write(&self, b: &bool) -> io::Result<()> {
        fs::write(&self.0, if *b { b"1" } else { b"0" })
    }
}

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

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

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

impl<P: AsRef<Path>> ReadAttr<u8> for SwappinessFile<P> {
    fn read(&self) -> io::Result<u8> {
        let file = self.0.as_ref();
        let s = fs::read_to_string(&file)?;
        match s.trim().parse() {
            Ok(u @ 0..=100) => Ok(u),
            Ok(_) => panic!("value of swappiness out of range"),
            Err(_) => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                format!("failed to parse swappiness in {}", file.display()),
            )),
        }
    }
}

impl<P: AsRef<Path>> WriteAttr<u8> for SwappinessFile<P> {
    fn write(&self, u: &u8) -> io::Result<()> {
        match *u {
            u @ 0..=100 => fs::write(&self.0, u.to_string()),
            _ => Err(io::Error::new(
                io::ErrorKind::InvalidData,
                "value of swappiness out of range",
            )),
        }
    }
}