liboj-cgroups 0.1.0

//! High-level APIs for cgroups v1.
//!
//! The APIs do not support all feature that cgroup can provide. If some features is not used when
//! implement an online judge or the kernel document doesn't provide enough information, the feature
//! will not be included in the APIs.
#[macro_use]
mod attr_file;
mod hierarchy;
pub mod subsystem;

#[cfg(test)]
mod tests;

use std::collections::HashSet;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};

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

use hierarchy::{Hierarchy, HierarchyNode};
use subsystem::{prelude::*, Subsystem};

lazy_static! {
    static ref SYSTEM_CGROUP_ROOT: &'static Path = Path::new("/sys/fs/cgroup/");
    static ref GLOBAL_CGROUP: GlobalContext =
        GlobalContext::new().expect("failed to initialize cgroup context");
}

struct GlobalContext;

impl GlobalContext {
    const GLOBAL_NAME: &'static str = "liboj";

    fn new() -> io::Result<GlobalContext> {
        let res = GlobalContext {};
        for controller in Subsystem::all() {
            fs::create_dir_all(res.subsystem_path(controller))?;
        }
        Ok(res)
    }

    fn subsystem_path(&self, sub: Subsystem) -> PathBuf {
        SYSTEM_CGROUP_ROOT.join(sub.name()).join(Self::GLOBAL_NAME)
    }

    fn subsystem(&self, sub: Subsystem) -> Hierarchy {
        Hierarchy::new(self.subsystem_path(sub))
    }
}

impl Drop for GlobalContext {
    fn drop(&mut self) {
        for sub in Subsystem::all() {
            let _ = fs::remove_dir_all(GLOBAL_CGROUP.subsystem_path(sub));
        }
    }
}

/// Cgroup context.
pub struct Context {
    name: String,
}

impl Context {
    fn get_controller(&self, sub: Subsystem) -> HierarchyNode {
        match GLOBAL_CGROUP.subsystem(sub).get_child(&self.name) {
            Some(controller) => controller,
            None => panic!(
                "Cgroup {} does not support {} subsystem",
                self.name,
                sub.name()
            ),
        }
    }

    pub fn add_process(&self, pid: Pid) -> io::Result<()> {
        for sub in Subsystem::all() {
            if let Some(controller) = GLOBAL_CGROUP.subsystem(sub).get_child(&self.name) {
                controller.procs().write(&pid)?;
            }
        }
        Ok(())
    }

    pub fn remove_process(&self, pid: Pid) -> io::Result<()> {
        for sub in Subsystem::all().map(|sub| GLOBAL_CGROUP.subsystem(sub).root()) {
            sub.procs().write(&pid)?;
        }
        Ok(())
    }

    pub fn cpu_controller(&self) -> Box<dyn '_ + CpuController> {
        Box::new(self.get_controller(Subsystem::Cpu))
    }

    pub fn cpuacct_controller(&self) -> Box<dyn '_ + CpuAcctController> {
        Box::new(self.get_controller(Subsystem::Cpuacct))
    }

    pub fn freezer_controller(&self) -> Box<dyn '_ + FreezerController> {
        Box::new(self.get_controller(Subsystem::Freezer))
    }

    pub fn memory_controller(&self) -> Box<dyn '_ + MemoryController> {
        Box::new(self.get_controller(Subsystem::Memory))
    }

    pub fn pids_controller(&self) -> Box<dyn '_ + PidsController> {
        Box::new(self.get_controller(Subsystem::Pids))
    }
}

pub struct Builder {
    name: Option<String>,
    subsystems: HashSet<Subsystem>,
}

impl Default for Builder {
    fn default() -> Builder {
        Builder::new().enable_all_subsystems()
    }
}

impl Builder {
    pub fn new() -> Builder {
        Builder {
            name: None,
            subsystems: HashSet::new(),
        }
    }

    pub fn set_name(mut self, name: String) -> Builder {
        self.name = Some(name);
        self
    }

    pub fn add_subsystem(mut self, sub: Subsystem) -> Builder {
        self.subsystems.insert(sub);
        self
    }

    pub fn enable_all_subsystems(mut self) -> Builder {
        for sub in Subsystem::all() {
            self.subsystems.insert(sub);
        }
        self
    }

    pub fn build(self) -> io::Result<Context> {
        use uuid::Uuid;
        let name = self.name.unwrap_or_else(|| format!("{:x}", Uuid::new_v4()));
        for controller in self.subsystems.into_iter() {
            GLOBAL_CGROUP.subsystem(controller).create(&name)?;
        }
        Ok(Context { name })
    }
}