1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

//! 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 })
    }
}