1use log::warn;
12use std::io::Write;
13use std::path::PathBuf;
14
15use crate::error::ErrorKind::*;
16use crate::error::*;
17use crate::{flat_keyed_to_vec, read_u64_from};
18
19use crate::{
20 ControllIdentifier, ControllerInternal, Controllers, HugePageResources, Resources, Subsystem,
21};
22
23#[derive(Debug, Clone)]
28pub struct HugeTlbController {
29 base: PathBuf,
30 path: PathBuf,
31 sizes: Vec<String>,
32 v2: bool,
33}
34
35impl ControllerInternal for HugeTlbController {
36 fn control_type(&self) -> Controllers {
37 Controllers::HugeTlb
38 }
39 fn get_path(&self) -> &PathBuf {
40 &self.path
41 }
42 fn get_path_mut(&mut self) -> &mut PathBuf {
43 &mut self.path
44 }
45 fn get_base(&self) -> &PathBuf {
46 &self.base
47 }
48
49 fn is_v2(&self) -> bool {
50 self.v2
51 }
52
53 fn apply(&self, res: &Resources) -> Result<()> {
54 let res: &HugePageResources = &res.hugepages;
56
57 for i in &res.limits {
58 let _ = self.set_limit_in_bytes(&i.size, i.limit);
59 if self.limit_in_bytes(&i.size)? != i.limit {
60 return Err(Error::new(Other));
61 }
62 }
63
64 Ok(())
65 }
66}
67
68impl ControllIdentifier for HugeTlbController {
69 fn controller_type() -> Controllers {
70 Controllers::HugeTlb
71 }
72}
73
74impl<'a> From<&'a Subsystem> for &'a HugeTlbController {
75 fn from(sub: &'a Subsystem) -> &'a HugeTlbController {
76 unsafe {
77 match sub {
78 Subsystem::HugeTlb(c) => c,
79 _ => {
80 assert_eq!(1, 0);
81 let v = std::mem::MaybeUninit::uninit();
82 v.assume_init()
83 }
84 }
85 }
86 }
87}
88
89impl HugeTlbController {
90 pub fn new(root: PathBuf, v2: bool) -> Self {
92 let sizes = get_hugepage_sizes();
93 Self {
94 base: root.clone(),
95 path: root,
96 sizes,
97 v2,
98 }
99 }
100
101 pub fn size_supported(&self, hugetlb_size: &str) -> bool {
103 for s in &self.sizes {
104 if s == hugetlb_size {
105 return true;
106 }
107 }
108 false
109 }
110
111 pub fn get_sizes(&self) -> Vec<String> {
112 self.sizes.clone()
113 }
114
115 fn failcnt_v2(&self, hugetlb_size: &str) -> Result<u64> {
116 self.open_path(&format!("hugetlb.{}.events", hugetlb_size), false)
117 .and_then(flat_keyed_to_vec)
118 .and_then(|x| {
119 if x.is_empty() {
120 return Err(Error::from_string(format!(
121 "get empty from hugetlb.{}.events",
122 hugetlb_size
123 )));
124 }
125 Ok(x[0].1 as u64)
126 })
127 }
128
129 pub fn failcnt(&self, hugetlb_size: &str) -> Result<u64> {
131 if self.v2 {
132 return self.failcnt_v2(hugetlb_size);
133 }
134 self.open_path(&format!("hugetlb.{}.failcnt", hugetlb_size), false)
135 .and_then(read_u64_from)
136 }
137
138 pub fn limit_in_bytes(&self, hugetlb_size: &str) -> Result<u64> {
141 let mut file_name = format!("hugetlb.{}.limit_in_bytes", hugetlb_size);
142 if self.v2 {
143 file_name = format!("hugetlb.{}.max", hugetlb_size);
144 }
145 self.open_path(&file_name, false).and_then(read_u64_from)
146 }
147
148 pub fn usage_in_bytes(&self, hugetlb_size: &str) -> Result<u64> {
151 let mut file = format!("hugetlb.{}.usage_in_bytes", hugetlb_size);
152 if self.v2 {
153 file = format!("hugetlb.{}.current", hugetlb_size);
154 }
155 self.open_path(&file, false).and_then(read_u64_from)
156 }
157
158 pub fn max_usage_in_bytes(&self, hugetlb_size: &str) -> Result<u64> {
161 self.open_path(
162 &format!("hugetlb.{}.max_usage_in_bytes", hugetlb_size),
163 false,
164 )
165 .and_then(read_u64_from)
166 }
167
168 pub fn set_limit_in_bytes(&self, hugetlb_size: &str, limit: u64) -> Result<()> {
171 let mut file_name = format!("hugetlb.{}.limit_in_bytes", hugetlb_size);
172 if self.v2 {
173 file_name = format!("hugetlb.{}.max", hugetlb_size);
174 }
175 self.open_path(&file_name, true).and_then(|mut file| {
176 file.write_all(limit.to_string().as_ref()).map_err(|e| {
177 Error::with_cause(WriteFailed(file_name.to_string(), limit.to_string()), e)
178 })
179 })
180 }
181}
182
183pub const HUGEPAGESIZE_DIR: &str = "/sys/kernel/mm/hugepages";
184use regex::Regex;
185use std::collections::HashMap;
186use std::fs;
187
188fn get_hugepage_sizes() -> Vec<String> {
189 let dirs = fs::read_dir(HUGEPAGESIZE_DIR);
190 if dirs.is_err() {
191 return Vec::new();
192 }
193
194 dirs.unwrap()
195 .filter_map(|e| {
196 let entry = e.map_err(|e| warn!("readdir error: {:?}", e)).ok()?;
197 let name = entry.file_name().into_string().unwrap();
198 let parts: Vec<&str> = name.split('-').collect();
199 if parts.len() != 2 {
200 return None;
201 }
202 let bmap = get_binary_size_map();
203 let size = parse_size(parts[1], &bmap)
204 .map_err(|e| warn!("parse_size error: {:?}", e))
205 .ok()?;
206 let dabbrs = get_decimal_abbrs();
207
208 Some(custom_size(size as f64, 1024.0, &dabbrs))
209 })
210 .collect()
211}
212
213pub const KB: u128 = 1000;
214pub const MB: u128 = 1000 * KB;
215pub const GB: u128 = 1000 * MB;
216pub const TB: u128 = 1000 * GB;
217pub const PB: u128 = 1000 * TB;
218
219#[allow(non_upper_case_globals)]
220pub const KiB: u128 = 1024;
221#[allow(non_upper_case_globals)]
222pub const MiB: u128 = 1024 * KiB;
223#[allow(non_upper_case_globals)]
224pub const GiB: u128 = 1024 * MiB;
225#[allow(non_upper_case_globals)]
226pub const TiB: u128 = 1024 * GiB;
227#[allow(non_upper_case_globals)]
228pub const PiB: u128 = 1024 * TiB;
229
230pub fn get_binary_size_map() -> HashMap<String, u128> {
231 let mut m = HashMap::new();
232 m.insert("k".to_string(), KiB);
233 m.insert("m".to_string(), MiB);
234 m.insert("g".to_string(), GiB);
235 m.insert("t".to_string(), TiB);
236 m.insert("p".to_string(), PiB);
237 m
238}
239
240pub fn get_decimal_size_map() -> HashMap<String, u128> {
241 let mut m = HashMap::new();
242 m.insert("k".to_string(), KB);
243 m.insert("m".to_string(), MB);
244 m.insert("g".to_string(), GB);
245 m.insert("t".to_string(), TB);
246 m.insert("p".to_string(), PB);
247 m
248}
249
250pub fn get_decimal_abbrs() -> Vec<String> {
251 let m = vec![
252 "B".to_string(),
253 "KB".to_string(),
254 "MB".to_string(),
255 "GB".to_string(),
256 "TB".to_string(),
257 "PB".to_string(),
258 "EB".to_string(),
259 "ZB".to_string(),
260 "YB".to_string(),
261 ];
262 m
263}
264
265fn parse_size(s: &str, m: &HashMap<String, u128>) -> Result<u128> {
266 let re = Regex::new(r"(?P<num>\d+)(?P<mul>[kKmMgGtTpP]?)[bB]?$");
267
268 if re.is_err() {
269 return Err(Error::new(InvalidBytesSize));
270 }
271 let caps = re.unwrap().captures(s).unwrap();
272
273 let num = caps.name("num");
274 let size: u128 = if let Some(num) = num {
275 let n = num.as_str().trim().parse::<u128>();
276 if n.is_err() {
277 return Err(Error::new(InvalidBytesSize));
278 }
279 n.unwrap()
280 } else {
281 return Err(Error::new(InvalidBytesSize));
282 };
283
284 let q = caps.name("mul");
285 let mul: u128 = if let Some(q) = q {
286 let t = m.get(q.as_str());
287 if let Some(t) = t {
288 *t
289 } else {
290 return Err(Error::new(InvalidBytesSize));
291 }
292 } else {
293 return Err(Error::new(InvalidBytesSize));
294 };
295
296 Ok(size * mul)
297}
298
299fn custom_size(mut size: f64, base: f64, m: &[String]) -> String {
300 let mut i = 0;
301 while size >= base && i < m.len() - 1 {
302 size /= base;
303 i += 1;
304 }
305
306 format!("{}{}", size, m[i].as_str())
307}