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
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
//! This module contains the implementation of the `cpu` cgroup subsystem.
//!
//! See the Kernel's documentation for more information about this subsystem, found at:
//!  [Documentation/scheduler/sched-design-CFS.txt](https://www.kernel.org/doc/Documentation/scheduler/sched-design-CFS.txt)
//!  paragraph 7 ("GROUP SCHEDULER EXTENSIONS TO CFS").
use std::fs::File;
use std::io::{Read, Write};
use std::path::PathBuf;

use error::*;
use error::ErrorKind::*;

use {
    ControllIdentifier, ControllerInternal, Controllers, CpuResources, Resources, Subsystem,
};

/// A controller that allows controlling the `cpu` subsystem of a Cgroup.
///
/// In essence, it allows gathering information about how much the tasks inside the control group
/// are using the CPU and creating rules that limit their usage. Note that this crate does not yet
/// support managing realtime tasks.
#[derive(Debug, Clone)]
pub struct CpuController {
    base: PathBuf,
    path: PathBuf,
}

/// The current state of the control group and its processes.
#[derive(Debug)]
pub struct Cpu {
    /// Reports CPU time statistics.
    ///
    /// Corresponds the `cpu.stat` file in `cpu` control group.
    pub stat: String,
}

impl ControllerInternal for CpuController {
    fn control_type(&self) -> Controllers {
        Controllers::Cpu
    }

    fn get_path(&self) -> &PathBuf {
        &self.path
    }

    fn get_path_mut(&mut self) -> &mut PathBuf {
        &mut self.path
    }

    fn get_base(&self) -> &PathBuf {
        &self.base
    }

    fn apply(&self, res: &Resources) -> Result<()> {
        // get the resources that apply to this controller
        let res: &CpuResources = &res.cpu;

        if res.update_values {
            // apply pid_max
            let _ = self.set_shares(res.shares);
            if self.shares()? != res.shares as u64 {
                return Err(Error::new(ErrorKind::Other));
            }

            let _ = self.set_cfs_period(res.period);
            if self.cfs_period()? != res.period as u64 {
                return Err(Error::new(ErrorKind::Other));
            }

            let _ = self.set_cfs_quota(res.quota as u64);
            if self.cfs_quota()? != res.quota as u64 {
                return Err(Error::new(ErrorKind::Other));
            }

            // TODO: rt properties (CONFIG_RT_GROUP_SCHED) are not yet supported
        }

        Ok(())
    }
}

impl ControllIdentifier for CpuController {
    fn controller_type() -> Controllers {
        Controllers::Cpu
    }
}

impl<'a> From<&'a Subsystem> for &'a CpuController {
    fn from(sub: &'a Subsystem) -> &'a CpuController {
        unsafe {
            match sub {
                Subsystem::Cpu(c) => c,
                _ => {
                    assert_eq!(1, 0);
                    ::std::mem::uninitialized()
                }
            }
        }
    }
}

fn read_u64_from(mut file: File) -> Result<u64> {
    let mut string = String::new();
    match file.read_to_string(&mut string) {
        Ok(_) => string.trim().parse().map_err(|e| Error::with_cause(ParseError, e)),
        Err(e) => Err(Error::with_cause(ReadFailed, e)),
    }
}

impl CpuController {
    /// Contructs a new `CpuController` with `oroot` serving as the root of the control group.
    pub fn new(oroot: PathBuf) -> Self {
        let mut root = oroot;
        root.push(Self::controller_type().to_string());
        Self {
            base: root.clone(),
            path: root,
        }
    }

    /// Returns CPU time statistics based on the processes in the control group.
    pub fn cpu(&self) -> Cpu {
        Cpu {
            stat: self
                .open_path("cpu.stat", false)
                .and_then(|mut file| {
                    let mut s = String::new();
                    let res = file.read_to_string(&mut s);
                    match res {
                        Ok(_) => Ok(s),
                        Err(e) => Err(Error::with_cause(ReadFailed, e)),
                    }
                }).unwrap_or("".to_string()),
        }
    }

    /// Configures the CPU bandwidth (in relative relation to other control groups and this control
    /// group's parent).
    ///
    /// For example, setting control group `A`'s `shares` to `100`, and control group `B`'s
    /// `shares` to `200` ensures that control group `B` receives twice as much as CPU bandwidth.
    /// (Assuming both `A` and `B` are of the same parent)
    pub fn set_shares(&self, shares: u64) -> Result<()> {
        self.open_path("cpu.shares", true).and_then(|mut file| {
            file.write_all(shares.to_string().as_ref())
                .map_err(|e| Error::with_cause(WriteFailed, e))
        })
    }

    /// Retrieve the CPU bandwidth that this control group (relative to other control groups and
    /// this control group's parent) can use.
    pub fn shares(&self) -> Result<u64> {
        self.open_path("cpu.shares", false).and_then(read_u64_from)
    }

    /// Specify a period (when using the CFS scheduler) of time in microseconds for how often this
    /// control group's access to the CPU should be reallocated.
    pub fn set_cfs_period(&self, us: u64) -> Result<()> {
        self.open_path("cpu.cfs_period_us", true)
            .and_then(|mut file| {
                file.write_all(us.to_string().as_ref())
                    .map_err(|e| Error::with_cause(WriteFailed, e))
            })
    }

    /// Retrieve the period of time of how often this cgroup's access to the CPU should be
    /// reallocated in microseconds.
    pub fn cfs_period(&self) -> Result<u64> {
        self.open_path("cpu.cfs_period_us", false)
            .and_then(read_u64_from)
    }

    /// Specify a quota (when using the CFS scheduler) of time in microseconds for which all tasks
    /// in this control group can run during one period (see: `set_cfs_period()`).
    pub fn set_cfs_quota(&self, us: u64) -> Result<()> {
        self.open_path("cpu.cfs_quota_us", true)
            .and_then(|mut file| {
                file.write_all(us.to_string().as_ref())
                    .map_err(|e| Error::with_cause(WriteFailed, e))
            })
    }

    /// Retrieve the quota of time for which all tasks in this cgroup can run during one period, in
    /// microseconds.
    pub fn cfs_quota(&self) -> Result<u64> {
        self.open_path("cpu.cfs_quota_us", false)
            .and_then(read_u64_from)
    }
}