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