1use std::fs::File;
13use std::io::{Read, Write};
14use std::path::PathBuf;
15
16use crate::error::ErrorKind::*;
17use crate::error::*;
18use crate::{parse_max_value, read_i64_from, read_u64_from};
19
20use crate::{
21 ControllIdentifier, ControllerInternal, Controllers, CpuResources, CustomizedAttribute,
22 MaxValue, Resources, Subsystem,
23};
24
25#[derive(Debug, Clone)]
31pub struct CpuController {
32 base: PathBuf,
33 path: PathBuf,
34 v2: bool,
35}
36
37#[derive(Debug)]
39#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
40pub struct Cpu {
41 pub stat: String,
45}
46
47#[derive(Debug)]
49struct CfsQuotaAndPeriod {
50 quota: MaxValue,
51 period: u64,
52}
53
54impl ControllerInternal for CpuController {
55 fn control_type(&self) -> Controllers {
56 Controllers::Cpu
57 }
58
59 fn get_path(&self) -> &PathBuf {
60 &self.path
61 }
62
63 fn get_path_mut(&mut self) -> &mut PathBuf {
64 &mut self.path
65 }
66
67 fn get_base(&self) -> &PathBuf {
68 &self.base
69 }
70
71 fn is_v2(&self) -> bool {
72 self.v2
73 }
74
75 fn apply(&self, res: &Resources) -> Result<()> {
76 let res: &CpuResources = &res.cpu;
78
79 update_and_test!(self, set_shares, res.shares, shares);
80 update_and_test!(self, set_cfs_period, res.period, cfs_period);
81 update_and_test!(self, set_cfs_quota, res.quota, cfs_quota);
82
83 res.attrs.iter().for_each(|(k, v)| {
84 let _ = self.set(k, v);
85 });
86
87 Ok(())
90 }
91}
92
93impl ControllIdentifier for CpuController {
94 fn controller_type() -> Controllers {
95 Controllers::Cpu
96 }
97}
98
99impl<'a> From<&'a Subsystem> for &'a CpuController {
100 fn from(sub: &'a Subsystem) -> &'a CpuController {
101 unsafe {
102 match sub {
103 Subsystem::Cpu(c) => c,
104 _ => {
105 assert_eq!(1, 0);
106 let v = std::mem::MaybeUninit::uninit();
107 v.assume_init()
108 }
109 }
110 }
111 }
112}
113
114impl CpuController {
115 pub fn new(root: PathBuf, v2: bool) -> Self {
117 Self {
118 base: root.clone(),
119 path: root,
120 v2,
121 }
122 }
123
124 pub fn cpu(&self) -> Cpu {
126 Cpu {
127 stat: self
128 .open_path("cpu.stat", false)
129 .and_then(|mut file| {
130 let mut s = String::new();
131 let res = file.read_to_string(&mut s);
132 match res {
133 Ok(_) => Ok(s),
134 Err(e) => Err(Error::with_cause(ReadFailed("cpu.stat".to_string()), e)),
135 }
136 })
137 .unwrap_or_default(),
138 }
139 }
140
141 pub fn set_shares(&self, shares: u64) -> Result<()> {
148 let mut file_name = "cpu.shares";
149 if self.v2 {
150 file_name = "cpu.weight";
151 }
152 self.open_path(file_name, true).and_then(|mut file| {
154 file.write_all(shares.to_string().as_ref()).map_err(|e| {
155 Error::with_cause(WriteFailed(file_name.to_string(), shares.to_string()), e)
156 })
157 })
158 }
159
160 pub fn shares(&self) -> Result<u64> {
163 let mut file = "cpu.shares";
164 if self.v2 {
165 file = "cpu.weight";
166 }
167 self.open_path(file, false).and_then(read_u64_from)
168 }
169
170 pub fn set_cfs_period(&self, us: u64) -> Result<()> {
173 if self.v2 {
174 return self.set_cfs_quota_and_period(None, Some(us));
175 }
176 self.open_path("cpu.cfs_period_us", true)
177 .and_then(|mut file| {
178 file.write_all(us.to_string().as_ref()).map_err(|e| {
179 Error::with_cause(
180 WriteFailed("cpu.cfs_period_us".to_string(), us.to_string()),
181 e,
182 )
183 })
184 })
185 }
186
187 pub fn cfs_period(&self) -> Result<u64> {
190 if self.v2 {
191 let current_value = self
192 .open_path("cpu.max", false)
193 .and_then(parse_cfs_quota_and_period)?;
194 return Ok(current_value.period);
195 }
196 self.open_path("cpu.cfs_period_us", false)
197 .and_then(read_u64_from)
198 }
199
200 pub fn set_cfs_quota(&self, us: i64) -> Result<()> {
203 if self.v2 {
204 return self.set_cfs_quota_and_period(Some(us), None);
205 }
206 self.open_path("cpu.cfs_quota_us", true)
207 .and_then(|mut file| {
208 file.write_all(us.to_string().as_ref()).map_err(|e| {
209 Error::with_cause(
210 WriteFailed("cpu.cfs_quota_us".to_string(), us.to_string()),
211 e,
212 )
213 })
214 })
215 }
216
217 pub fn cfs_quota(&self) -> Result<i64> {
220 if self.v2 {
221 let current_value = self
222 .open_path("cpu.max", false)
223 .and_then(parse_cfs_quota_and_period)?;
224 return Ok(current_value.quota.to_i64());
225 }
226
227 self.open_path("cpu.cfs_quota_us", false)
228 .and_then(read_i64_from)
229 }
230
231 pub fn set_cfs_quota_and_period(&self, quota: Option<i64>, period: Option<u64>) -> Result<()> {
232 if !self.v2 {
233 if let Some(q) = quota {
234 self.set_cfs_quota(q)?;
235 }
236 if let Some(p) = period {
237 self.set_cfs_period(p)?;
238 }
239 return Ok(());
240 }
241
242 let current_value = self
252 .open_path("cpu.max", false)
253 .and_then(parse_cfs_quota_and_period)?;
254
255 let new_quota = if let Some(q) = quota {
256 if q > 0 {
257 q.to_string()
258 } else {
259 "max".to_string()
260 }
261 } else {
262 current_value.quota.to_string()
263 };
264
265 let new_period = if let Some(p) = period {
266 p.to_string()
267 } else {
268 current_value.period.to_string()
269 };
270
271 let line = format!("{} {}", new_quota, new_period);
272 self.open_path("cpu.max", true).and_then(|mut file| {
273 file.write_all(line.as_ref())
274 .map_err(|e| Error::with_cause(WriteFailed("cpu.max".to_string(), line), e))
275 })
276 }
277
278 pub fn set_rt_runtime(&self, us: i64) -> Result<()> {
279 self.open_path("cpu.rt_runtime_us", true)
280 .and_then(|mut file| {
281 file.write_all(us.to_string().as_ref()).map_err(|e| {
282 Error::with_cause(
283 WriteFailed("cpu.rt_runtime_us".to_string(), us.to_string()),
284 e,
285 )
286 })
287 })
288 }
289
290 pub fn set_rt_period_us(&self, us: u64) -> Result<()> {
291 self.open_path("cpu.rt_period_us", true)
292 .and_then(|mut file| {
293 file.write_all(us.to_string().as_ref()).map_err(|e| {
294 Error::with_cause(
295 WriteFailed("cpu.rt_period_us".to_string(), us.to_string()),
296 e,
297 )
298 })
299 })
300 }
301}
302
303impl CustomizedAttribute for CpuController {}
304
305fn parse_cfs_quota_and_period(mut file: File) -> Result<CfsQuotaAndPeriod> {
306 let mut content = String::new();
307 file.read_to_string(&mut content)
308 .map_err(|e| Error::with_cause(ReadFailed("cpu.max".to_string()), e))?;
309
310 let fields = content.trim().split(' ').collect::<Vec<&str>>();
311 if fields.len() != 2 {
312 return Err(Error::from_string(format!("invaild format: {}", content)));
313 }
314
315 let quota = parse_max_value(fields[0])?;
316 let period = fields[1]
317 .parse::<u64>()
318 .map_err(|e| Error::with_cause(ParseError, e))?;
319
320 Ok(CfsQuotaAndPeriod { quota, period })
321}