1use std::collections::HashMap;
2use std::fs::File;
3use std::io::{BufRead, BufReader, Read};
4use std::mem;
5use std::path::{Path, PathBuf};
6use std::sync::atomic::{AtomicUsize, Ordering};
7use std::sync::Once;
8
9use libc;
10
11macro_rules! debug {
12 ($($args:expr),*) => ({
13 if false {
14 //if true {
15 println!($($args),*);
16 }
17 });
18}
19
20macro_rules! some {
21 ($e:expr) => {{
22 match $e {
23 Some(v) => v,
24 None => {
25 debug!("NONE: {:?}", stringify!($e));
26 return None;
27 }
28 }
29 }};
30}
31
32pub fn get_num_cpus() -> usize {
33 match cgroups_num_cpus() {
34 Some(n: usize) => n,
35 None => logical_cpus(),
36 }
37}
38
39fn logical_cpus() -> usize {
40 let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
41 if unsafe { libc::sched_getaffinity(pid:0, cpusetsize:mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
42 let mut count: u32 = 0;
43 for i: usize in 0..libc::CPU_SETSIZE as usize {
44 if unsafe { libc::CPU_ISSET(cpu:i, &set) } {
45 count += 1
46 }
47 }
48 count as usize
49 } else {
50 let cpus: i64 = unsafe { libc::sysconf(name:libc::_SC_NPROCESSORS_ONLN) };
51 if cpus < 1 {
52 1
53 } else {
54 cpus as usize
55 }
56 }
57}
58
59pub fn get_num_physical_cpus() -> usize {
60 let file = match File::open("/proc/cpuinfo") {
61 Ok(val) => val,
62 Err(_) => return get_num_cpus(),
63 };
64 let reader = BufReader::new(file);
65 let mut map = HashMap::new();
66 let mut physid: u32 = 0;
67 let mut cores: usize = 0;
68 let mut chgcount = 0;
69 for line in reader.lines().filter_map(|result| result.ok()) {
70 let mut it = line.split(':');
71 let (key, value) = match (it.next(), it.next()) {
72 (Some(key), Some(value)) => (key.trim(), value.trim()),
73 _ => continue,
74 };
75 if key == "physical id" {
76 match value.parse() {
77 Ok(val) => physid = val,
78 Err(_) => break,
79 };
80 chgcount += 1;
81 }
82 if key == "cpu cores" {
83 match value.parse() {
84 Ok(val) => cores = val,
85 Err(_) => break,
86 };
87 chgcount += 1;
88 }
89 if chgcount == 2 {
90 map.insert(physid, cores);
91 chgcount = 0;
92 }
93 }
94 let count = map.into_iter().fold(0, |acc, (_, cores)| acc + cores);
95
96 if count == 0 {
97 get_num_cpus()
98 } else {
99 count
100 }
101}
102
103/// Cached CPUs calculated from cgroups.
104///
105/// If 0, check logical cpus.
106// Allow deprecation warnings, we want to work on older rustc
107#[allow(warnings)]
108static CGROUPS_CPUS: AtomicUsize = ::std::sync::atomic::ATOMIC_USIZE_INIT;
109
110fn cgroups_num_cpus() -> Option<usize> {
111 #[allow(warnings)]
112 static ONCE: Once = ::std::sync::ONCE_INIT;
113
114 ONCE.call_once(init_cgroups);
115
116 let cpus: usize = CGROUPS_CPUS.load(order:Ordering::Acquire);
117
118 if cpus > 0 {
119 Some(cpus)
120 } else {
121 None
122 }
123}
124
125fn init_cgroups() {
126 // Should only be called once
127 debug_assert!(CGROUPS_CPUS.load(Ordering::SeqCst) == 0);
128
129 // Fails in Miri by default (cannot open files), and Miri does not have parallelism anyway.
130 if cfg!(miri) {
131 return;
132 }
133
134 if let Some(quota: usize) = load_cgroups(cgroup_proc:"/proc/self/cgroup", mountinfo_proc:"/proc/self/mountinfo") {
135 if quota == 0 {
136 return;
137 }
138
139 let logical: usize = logical_cpus();
140 let count: usize = ::std::cmp::min(v1:quota, v2:logical);
141
142 CGROUPS_CPUS.store(val:count, order:Ordering::SeqCst);
143 }
144}
145
146fn load_cgroups<P1, P2>(cgroup_proc: P1, mountinfo_proc: P2) -> Option<usize>
147where
148 P1: AsRef<Path>,
149 P2: AsRef<Path>,
150{
151 let subsys: Subsys = some!(Subsys::load_cpu(cgroup_proc));
152 let mntinfo: MountInfo = some!(MountInfo::load_cpu(mountinfo_proc, subsys.version));
153 let cgroup: Cgroup = some!(Cgroup::translate(mntinfo, subsys));
154 cgroup.cpu_quota()
155}
156
157#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158enum CgroupVersion {
159 V1,
160 V2,
161}
162
163struct Cgroup {
164 version: CgroupVersion,
165 base: PathBuf,
166}
167
168struct MountInfo {
169 version: CgroupVersion,
170 root: String,
171 mount_point: String,
172}
173
174struct Subsys {
175 version: CgroupVersion,
176 base: String,
177}
178
179impl Cgroup {
180 fn new(version: CgroupVersion, dir: PathBuf) -> Cgroup {
181 Cgroup { version: version, base: dir }
182 }
183
184 fn translate(mntinfo: MountInfo, subsys: Subsys) -> Option<Cgroup> {
185 // Translate the subsystem directory via the host paths.
186 debug!(
187 "subsys = {:?}; root = {:?}; mount_point = {:?}",
188 subsys.base, mntinfo.root, mntinfo.mount_point
189 );
190
191 let rel_from_root = some!(Path::new(&subsys.base).strip_prefix(&mntinfo.root).ok());
192
193 debug!("rel_from_root: {:?}", rel_from_root);
194
195 // join(mp.MountPoint, relPath)
196 let mut path = PathBuf::from(mntinfo.mount_point);
197 path.push(rel_from_root);
198 Some(Cgroup::new(mntinfo.version, path))
199 }
200
201 fn cpu_quota(&self) -> Option<usize> {
202 let (quota_us, period_us) = match self.version {
203 CgroupVersion::V1 => (some!(self.quota_us()), some!(self.period_us())),
204 CgroupVersion::V2 => some!(self.max()),
205 };
206
207 // protect against dividing by zero
208 if period_us == 0 {
209 return None;
210 }
211
212 // Ceil the division, since we want to be able to saturate
213 // the available CPUs, and flooring would leave a CPU un-utilized.
214
215 Some((quota_us as f64 / period_us as f64).ceil() as usize)
216 }
217
218 fn quota_us(&self) -> Option<usize> {
219 self.param("cpu.cfs_quota_us")
220 }
221
222 fn period_us(&self) -> Option<usize> {
223 self.param("cpu.cfs_period_us")
224 }
225
226 fn max(&self) -> Option<(usize, usize)> {
227 let max = some!(self.raw_param("cpu.max"));
228 let mut max = some!(max.lines().next()).split(' ');
229
230 let quota = some!(max.next().and_then(|quota| quota.parse().ok()));
231 let period = some!(max.next().and_then(|period| period.parse().ok()));
232
233 Some((quota, period))
234 }
235
236 fn param(&self, param: &str) -> Option<usize> {
237 let buf = some!(self.raw_param(param));
238
239 buf.trim().parse().ok()
240 }
241
242 fn raw_param(&self, param: &str) -> Option<String> {
243 let mut file = some!(File::open(self.base.join(param)).ok());
244
245 let mut buf = String::new();
246 some!(file.read_to_string(&mut buf).ok());
247
248 Some(buf)
249 }
250}
251
252impl MountInfo {
253 fn load_cpu<P: AsRef<Path>>(proc_path: P, version: CgroupVersion) -> Option<MountInfo> {
254 let file = some!(File::open(proc_path).ok());
255 let file = BufReader::new(file);
256
257 file.lines()
258 .filter_map(|result| result.ok())
259 .filter_map(MountInfo::parse_line)
260 .find(|mount_info| mount_info.version == version)
261 }
262
263 fn parse_line(line: String) -> Option<MountInfo> {
264 let mut fields = line.split(' ');
265
266 // 7 5 0:6 </> /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct
267 let mnt_root = some!(fields.nth(3));
268 // 7 5 0:6 / </sys/fs/cgroup/cpu,cpuacct> rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct
269 let mnt_point = some!(fields.next());
270
271 // Ignore all fields until the separator(-).
272 // Note: there could be zero or more optional fields before hyphen.
273 // See: https://man7.org/linux/man-pages/man5/proc.5.html
274 // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 <-> cgroup cgroup rw,cpu,cpuacct
275 // Note: we cannot use `?` here because we need to support Rust 1.13.
276 match fields.find(|&s| s == "-") {
277 Some(_) => {}
278 None => return None,
279 };
280
281 // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - <cgroup> cgroup rw,cpu,cpuacct
282 let version = match fields.next() {
283 Some("cgroup") => CgroupVersion::V1,
284 Some("cgroup2") => CgroupVersion::V2,
285 _ => return None,
286 };
287
288 // cgroups2 only has a single mount point
289 if version == CgroupVersion::V1 {
290 // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup <rw,cpu,cpuacct>
291 let super_opts = some!(fields.nth(1));
292
293 // We only care about the 'cpu' option
294 if !super_opts.split(',').any(|opt| opt == "cpu") {
295 return None;
296 }
297 }
298
299 Some(MountInfo {
300 version: version,
301 root: mnt_root.to_owned(),
302 mount_point: mnt_point.to_owned(),
303 })
304 }
305}
306
307impl Subsys {
308 fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<Subsys> {
309 let file = some!(File::open(proc_path).ok());
310 let file = BufReader::new(file);
311
312 file.lines()
313 .filter_map(|result| result.ok())
314 .filter_map(Subsys::parse_line)
315 .fold(None, |previous, line| {
316 // already-found v1 trumps v2 since it explicitly specifies its controllers
317 if previous.is_some() && line.version == CgroupVersion::V2 {
318 return previous;
319 }
320
321 Some(line)
322 })
323 }
324
325 fn parse_line(line: String) -> Option<Subsys> {
326 // Example format:
327 // 11:cpu,cpuacct:/
328 let mut fields = line.split(':');
329
330 let sub_systems = some!(fields.nth(1));
331
332 let version = if sub_systems.is_empty() {
333 CgroupVersion::V2
334 } else {
335 CgroupVersion::V1
336 };
337
338 if version == CgroupVersion::V1 && !sub_systems.split(',').any(|sub| sub == "cpu") {
339 return None;
340 }
341
342 fields.next().map(|path| Subsys {
343 version: version,
344 base: path.to_owned(),
345 })
346 }
347}
348
349#[cfg(test)]
350mod tests {
351 mod v1 {
352 use super::super::{Cgroup, CgroupVersion, MountInfo, Subsys};
353 use std::path::{Path, PathBuf};
354
355 // `static_in_const` feature is not stable in Rust 1.13.
356 static FIXTURES_PROC: &'static str = "fixtures/cgroups/proc/cgroups";
357
358 static FIXTURES_CGROUPS: &'static str = "fixtures/cgroups/cgroups";
359
360 macro_rules! join {
361 ($base:expr, $($path:expr),+) => ({
362 Path::new($base)
363 $(.join($path))+
364 })
365 }
366
367 #[test]
368 fn test_load_mountinfo() {
369 // test only one optional fields
370 let path = join!(FIXTURES_PROC, "mountinfo");
371
372 let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V1).unwrap();
373
374 assert_eq!(mnt_info.root, "/");
375 assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
376
377 // test zero optional field
378 let path = join!(FIXTURES_PROC, "mountinfo_zero_opt");
379
380 let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V1).unwrap();
381
382 assert_eq!(mnt_info.root, "/");
383 assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
384
385 // test multi optional fields
386 let path = join!(FIXTURES_PROC, "mountinfo_multi_opt");
387
388 let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V1).unwrap();
389
390 assert_eq!(mnt_info.root, "/");
391 assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
392 }
393
394 #[test]
395 fn test_load_subsys() {
396 let path = join!(FIXTURES_PROC, "cgroup");
397
398 let subsys = Subsys::load_cpu(path).unwrap();
399
400 assert_eq!(subsys.base, "/");
401 assert_eq!(subsys.version, CgroupVersion::V1);
402 }
403
404 #[test]
405 fn test_cgroup_mount() {
406 let cases = &[
407 ("/", "/sys/fs/cgroup/cpu", "/", Some("/sys/fs/cgroup/cpu")),
408 (
409 "/docker/01abcd",
410 "/sys/fs/cgroup/cpu",
411 "/docker/01abcd",
412 Some("/sys/fs/cgroup/cpu"),
413 ),
414 (
415 "/docker/01abcd",
416 "/sys/fs/cgroup/cpu",
417 "/docker/01abcd/",
418 Some("/sys/fs/cgroup/cpu"),
419 ),
420 (
421 "/docker/01abcd",
422 "/sys/fs/cgroup/cpu",
423 "/docker/01abcd/large",
424 Some("/sys/fs/cgroup/cpu/large"),
425 ),
426 // fails
427 ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/", None),
428 ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/docker", None),
429 ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/elsewhere", None),
430 (
431 "/docker/01abcd",
432 "/sys/fs/cgroup/cpu",
433 "/docker/01abcd-other-dir",
434 None,
435 ),
436 ];
437
438 for &(root, mount_point, subsys, expected) in cases.iter() {
439 let mnt_info = MountInfo {
440 version: CgroupVersion::V1,
441 root: root.into(),
442 mount_point: mount_point.into(),
443 };
444 let subsys = Subsys {
445 version: CgroupVersion::V1,
446 base: subsys.into(),
447 };
448
449 let actual = Cgroup::translate(mnt_info, subsys).map(|c| c.base);
450 let expected = expected.map(PathBuf::from);
451 assert_eq!(actual, expected);
452 }
453 }
454
455 #[test]
456 fn test_cgroup_cpu_quota() {
457 let cgroup = Cgroup::new(CgroupVersion::V1, join!(FIXTURES_CGROUPS, "good"));
458 assert_eq!(cgroup.cpu_quota(), Some(6));
459 }
460
461 #[test]
462 fn test_cgroup_cpu_quota_divide_by_zero() {
463 let cgroup = Cgroup::new(CgroupVersion::V1, join!(FIXTURES_CGROUPS, "zero-period"));
464 assert!(cgroup.quota_us().is_some());
465 assert_eq!(cgroup.period_us(), Some(0));
466 assert_eq!(cgroup.cpu_quota(), None);
467 }
468
469 #[test]
470 fn test_cgroup_cpu_quota_ceil() {
471 let cgroup = Cgroup::new(CgroupVersion::V1, join!(FIXTURES_CGROUPS, "ceil"));
472 assert_eq!(cgroup.cpu_quota(), Some(2));
473 }
474 }
475
476 mod v2 {
477 use super::super::{Cgroup, CgroupVersion, MountInfo, Subsys};
478 use std::path::{Path, PathBuf};
479
480 // `static_in_const` feature is not stable in Rust 1.13.
481 static FIXTURES_PROC: &'static str = "fixtures/cgroups2/proc/cgroups";
482
483 static FIXTURES_CGROUPS: &'static str = "fixtures/cgroups2/cgroups";
484
485 macro_rules! join {
486 ($base:expr, $($path:expr),+) => ({
487 Path::new($base)
488 $(.join($path))+
489 })
490 }
491
492 #[test]
493 fn test_load_mountinfo() {
494 // test only one optional fields
495 let path = join!(FIXTURES_PROC, "mountinfo");
496
497 let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V2).unwrap();
498
499 assert_eq!(mnt_info.root, "/");
500 assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup");
501 }
502
503 #[test]
504 fn test_load_subsys() {
505 let path = join!(FIXTURES_PROC, "cgroup");
506
507 let subsys = Subsys::load_cpu(path).unwrap();
508
509 assert_eq!(subsys.base, "/");
510 assert_eq!(subsys.version, CgroupVersion::V2);
511 }
512
513 #[test]
514 fn test_load_subsys_multi() {
515 let path = join!(FIXTURES_PROC, "cgroup_multi");
516
517 let subsys = Subsys::load_cpu(path).unwrap();
518
519 assert_eq!(subsys.base, "/");
520 assert_eq!(subsys.version, CgroupVersion::V1);
521 }
522
523 #[test]
524 fn test_cgroup_mount() {
525 let cases = &[
526 ("/", "/sys/fs/cgroup/cpu", "/", Some("/sys/fs/cgroup/cpu")),
527 (
528 "/docker/01abcd",
529 "/sys/fs/cgroup/cpu",
530 "/docker/01abcd",
531 Some("/sys/fs/cgroup/cpu"),
532 ),
533 (
534 "/docker/01abcd",
535 "/sys/fs/cgroup/cpu",
536 "/docker/01abcd/",
537 Some("/sys/fs/cgroup/cpu"),
538 ),
539 (
540 "/docker/01abcd",
541 "/sys/fs/cgroup/cpu",
542 "/docker/01abcd/large",
543 Some("/sys/fs/cgroup/cpu/large"),
544 ),
545 // fails
546 ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/", None),
547 ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/docker", None),
548 ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/elsewhere", None),
549 (
550 "/docker/01abcd",
551 "/sys/fs/cgroup/cpu",
552 "/docker/01abcd-other-dir",
553 None,
554 ),
555 ];
556
557 for &(root, mount_point, subsys, expected) in cases.iter() {
558 let mnt_info = MountInfo {
559 version: CgroupVersion::V1,
560 root: root.into(),
561 mount_point: mount_point.into(),
562 };
563 let subsys = Subsys {
564 version: CgroupVersion::V1,
565 base: subsys.into(),
566 };
567
568 let actual = Cgroup::translate(mnt_info, subsys).map(|c| c.base);
569 let expected = expected.map(PathBuf::from);
570 assert_eq!(actual, expected);
571 }
572 }
573
574 #[test]
575 fn test_cgroup_cpu_quota() {
576 let cgroup = Cgroup::new(CgroupVersion::V2, join!(FIXTURES_CGROUPS, "good"));
577 assert_eq!(cgroup.cpu_quota(), Some(6));
578 }
579
580 #[test]
581 fn test_cgroup_cpu_quota_divide_by_zero() {
582 let cgroup = Cgroup::new(CgroupVersion::V2, join!(FIXTURES_CGROUPS, "zero-period"));
583 let period = cgroup.max().map(|max| max.1);
584
585 assert_eq!(period, Some(0));
586 assert_eq!(cgroup.cpu_quota(), None);
587 }
588
589 #[test]
590 fn test_cgroup_cpu_quota_ceil() {
591 let cgroup = Cgroup::new(CgroupVersion::V2, join!(FIXTURES_CGROUPS, "ceil"));
592 assert_eq!(cgroup.cpu_quota(), Some(2));
593 }
594 }
595}
596