liboj-cgroups 0.1.0

Cgroup wrapping for liboj
Documentation
pub mod cpu;
pub mod cpuacct;
pub mod freezer;
pub mod memory;
pub mod pids;
pub mod prelude;

use std::fs;
use std::io;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};

use nix::unistd::Pid;

use crate::{
    attr_file::{AttrFile, ReadAttr, WriteAttr},
    hierarchy::HierarchyNode,
};

pub trait Controller {
    /// List of thread group IDs in the cgroup. This list is not guaranteed to be sorted
    /// or free of duplicate TGIDs, and userspace should sort/uniquify the list if this
    /// property is required. Writing a thread group ID into this file moves all threads
    /// in that group into this cgroup.
    fn procs(&self) -> Box<dyn '_ + AttrFile<Vec<Pid>, Pid>>;
    /// List of tasks (by PID) attached to that cgroup. This list is not guaranteed to be
    /// sorted. Writing a thread ID into this file moves the thread into this cgroup.
    fn tasks(&self) -> Box<dyn '_ + AttrFile<Vec<Pid>, Pid>>;
    /// Flag if run the release agent on exit?
    fn notify_on_release(&self) -> Box<dyn '_ + AttrFile<bool>>;
    /// The path to use for release notifications (exists in the top cgroup only)
    fn release_agent(&self) -> Box<dyn '_ + AttrFile<PathBuf, Path>>;
}

/// Hierarchy in the cgroup.
impl Controller for HierarchyNode {
    fn procs(&self) -> Box<dyn '_ + AttrFile<Vec<Pid>, Pid>> {
        let file = self.as_path().join("cgroup.procs");
        Box::new(PidFile(file))
    }

    fn tasks(&self) -> Box<dyn '_ + AttrFile<Vec<Pid>, Pid>> {
        let file = self.as_path().join("tasks");
        Box::new(PidFile(file))
    }

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

    fn release_agent(&self) -> Box<dyn '_ + AttrFile<PathBuf, Path>> {
        let file = self.as_path().join("release_agent");
        Box::new(ReleaseAgentFile(file))
    }
}

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

impl<P: AsRef<Path>> ReadAttr<Vec<Pid>> for PidFile<P> {
    fn read(&self) -> io::Result<Vec<Pid>> {
        let file = self.0.as_ref();
        let s = fs::read_to_string(&file)?;
        let mut pids: Vec<Pid> = Vec::new();
        for pid in s.split_whitespace() {
            match pid.trim().parse() {
                Ok(pid) => pids.push(Pid::from_raw(pid)),
                Err(_) => {
                    return Err(io::Error::new(
                        io::ErrorKind::InvalidData,
                        format!("failed to get pid from {}", file.display()),
                    ));
                }
            }
        }
        Ok(pids)
    }
}

impl<P: AsRef<Path>> WriteAttr<Pid> for PidFile<P> {
    fn write(&self, pid: &Pid) -> io::Result<()> {
        fs::write(&self.0, pid.to_string())
    }
}

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

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

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

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

impl<P: AsRef<Path>> ReadAttr<PathBuf> for ReleaseAgentFile<P> {
    fn read(&self) -> io::Result<PathBuf> {
        let file = self.0.as_ref();
        let s = fs::read_to_string(&file)?;
        Ok(PathBuf::from(s))
    }
}

impl<P: AsRef<Path>> WriteAttr<Path> for ReleaseAgentFile<P> {
    fn write(&self, p: &Path) -> io::Result<()> {
        fs::write(&self.0, &p.as_os_str().as_bytes())
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Subsystem {
    Cpu,
    Cpuacct,
    Freezer,
    Memory,
    Pids,
}

impl Subsystem {
    pub fn all() -> impl Iterator<Item = Subsystem> {
        use Subsystem::*;
        vec![Cpu, Cpuacct, Freezer, Memory, Pids].into_iter()
    }

    pub fn name(self) -> &'static str {
        use Subsystem::*;
        match self {
            Cpu => "cpu",
            Cpuacct => "cpuacct",
            Freezer => "freezer",
            Memory => "memory",
            Pids => "pids",
        }
    }
}