/rust/registry/src/index.crates.io-1949cf8c6b5b557f/num_cpus-1.17.0/src/linux.rs

Source
use std::collections::HashMap;
use std::fs::File;
use std::io::{BufRead, BufReader, Read};
use std::mem;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Once;

use libc;

macro_rules! debug {
    ($($args:expr),*) => ({
        if false {
        //if true {
            println!($($args),*);
        }
    });
}

macro_rules! some {
    ($e:expr) => {{
        match $e {
            Some(v) => v,
            None => {
                debug!("NONE: {:?}", stringify!($e));
                return None;
            }
        }
    }};
}

pub fn get_num_cpus() -> usize {
    match cgroups_num_cpus() {
        Some(n) => n,
        None => logical_cpus(),
    }
}

fn logical_cpus() -> usize {
    let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
    if unsafe { libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
        let mut count: u32 = 0;
        for i in 0..libc::CPU_SETSIZE as usize {
            if unsafe { libc::CPU_ISSET(i, &set) } {
                count += 1
            }
        }
        count as usize
    } else {
        let cpus = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) };
        if cpus < 1 {
            1
        } else {
            cpus as usize
        }
    }
}

pub fn get_num_physical_cpus() -> usize {
    let file = match File::open("/proc/cpuinfo") {
        Ok(val) => val,
        Err(_) => return get_num_cpus(),
    };
    let reader = BufReader::new(file);
    let mut map = HashMap::new();
    let mut physid: u32 = 0;
    let mut cores: usize = 0;
    let mut chgcount = 0;
    for line in reader.lines().filter_map(|result| result.ok()) {
        let mut it = line.split(':');
        let (key, value) = match (it.next(), it.next()) {
            (Some(key), Some(value)) => (key.trim(), value.trim()),
            _ => continue,
        };
        if key == "physical id" {
            match value.parse() {
                Ok(val) => physid = val,
                Err(_) => break,
            };
            chgcount += 1;
        }
        if key == "cpu cores" {
            match value.parse() {
                Ok(val) => cores = val,
                Err(_) => break,
            };
            chgcount += 1;
        }
        if chgcount == 2 {
            map.insert(physid, cores);
            chgcount = 0;
        }
    }
    let count = map.into_iter().fold(0, |acc, (_, cores)| acc + cores);

    if count == 0 {
        get_num_cpus()
    } else {
        count
    }
}

/// Cached CPUs calculated from cgroups.
///
/// If 0, check logical cpus.
// Allow deprecation warnings, we want to work on older rustc
#[allow(warnings)]
static CGROUPS_CPUS: AtomicUsize = ::std::sync::atomic::ATOMIC_USIZE_INIT;

fn cgroups_num_cpus() -> Option<usize> {
    #[allow(warnings)]
    static ONCE: Once = ::std::sync::ONCE_INIT;

    ONCE.call_once(init_cgroups);

    let cpus = CGROUPS_CPUS.load(Ordering::Acquire);

    if cpus > 0 {
        Some(cpus)
    } else {
        None
    }
}

fn init_cgroups() {
    // Should only be called once
    debug_assert!(CGROUPS_CPUS.load(Ordering::SeqCst) == 0);

    // Fails in Miri by default (cannot open files), and Miri does not have parallelism anyway.
    if cfg!(miri) {
        return;
    }

    if let Some(quota) = load_cgroups("/proc/self/cgroup", "/proc/self/mountinfo") {
        if quota == 0 {
            return;
        }

        let logical = logical_cpus();
        let count = ::std::cmp::min(quota, logical);

        CGROUPS_CPUS.store(count, Ordering::SeqCst);
    }
}

fn load_cgroups<P1, P2>(cgroup_proc: P1, mountinfo_proc: P2) -> Option<usize>
where
    P1: AsRef<Path>,
    P2: AsRef<Path>,
{
    let subsys = some!(Subsys::load_cpu(cgroup_proc));
    let mntinfo = some!(MountInfo::load_cpu(mountinfo_proc, subsys.version));
    let cgroup = some!(Cgroup::translate(mntinfo, subsys));
    cgroup.cpu_quota()
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum CgroupVersion {
    V1,
    V2,
}

struct Cgroup {
    version: CgroupVersion,
    base: PathBuf,
}

struct MountInfo {
    version: CgroupVersion,
    root: String,
    mount_point: String,
}

struct Subsys {
    version: CgroupVersion,
    base: String,
}

impl Cgroup {
    fn new(version: CgroupVersion, dir: PathBuf) -> Cgroup {
        Cgroup { version: version, base: dir }
    }

    fn translate(mntinfo: MountInfo, subsys: Subsys) -> Option<Cgroup> {
        // Translate the subsystem directory via the host paths.
        debug!(
            "subsys = {:?}; root = {:?}; mount_point = {:?}",
            subsys.base, mntinfo.root, mntinfo.mount_point
        );

        let rel_from_root = some!(Path::new(&subsys.base).strip_prefix(&mntinfo.root).ok());

        debug!("rel_from_root: {:?}", rel_from_root);

        // join(mp.MountPoint, relPath)
        let mut path = PathBuf::from(mntinfo.mount_point);
        path.push(rel_from_root);
        Some(Cgroup::new(mntinfo.version, path))
    }

    fn cpu_quota(&self) -> Option<usize> {
        let (quota_us, period_us) = match self.version {
            CgroupVersion::V1 => (some!(self.quota_us()), some!(self.period_us())),
            CgroupVersion::V2 => some!(self.max()),
        };

        // protect against dividing by zero
        if period_us == 0 {
            return None;
        }

        // Ceil the division, since we want to be able to saturate
        // the available CPUs, and flooring would leave a CPU un-utilized.

        Some((quota_us as f64 / period_us as f64).ceil() as usize)
    }

    fn quota_us(&self) -> Option<usize> {
        self.param("cpu.cfs_quota_us")
    }

    fn period_us(&self) -> Option<usize> {
        self.param("cpu.cfs_period_us")
    }

    fn max(&self) -> Option<(usize, usize)> {
        let max = some!(self.raw_param("cpu.max"));
        let mut max = some!(max.lines().next()).split(' ');

        let quota = some!(max.next().and_then(|quota| quota.parse().ok()));
        let period = some!(max.next().and_then(|period| period.parse().ok()));

        Some((quota, period))
    }

    fn param(&self, param: &str) -> Option<usize> {
        let buf = some!(self.raw_param(param));

        buf.trim().parse().ok()
    }

    fn raw_param(&self, param: &str) -> Option<String> {
        let mut file = some!(File::open(self.base.join(param)).ok());

        let mut buf = String::new();
        some!(file.read_to_string(&mut buf).ok());

        Some(buf)
    }
}

impl MountInfo {
    fn load_cpu<P: AsRef<Path>>(proc_path: P, version: CgroupVersion) -> Option<MountInfo> {
        let file = some!(File::open(proc_path).ok());
        let file = BufReader::new(file);

        file.lines()
            .filter_map(|result| result.ok())
            .filter_map(MountInfo::parse_line)
            .find(|mount_info| mount_info.version == version)
    }

    fn parse_line(line: String) -> Option<MountInfo> {
        let mut fields = line.split(' ');

        // 7 5 0:6 </> /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct
        let mnt_root = some!(fields.nth(3));
        // 7 5 0:6 / </sys/fs/cgroup/cpu,cpuacct> rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup rw,cpu,cpuacct
        let mnt_point = some!(fields.next());

        // Ignore all fields until the separator(-).
        // Note: there could be zero or more optional fields before hyphen.
        // See: https://man7.org/linux/man-pages/man5/proc.5.html
        // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 <-> cgroup cgroup rw,cpu,cpuacct
        // Note: we cannot use `?` here because we need to support Rust 1.13.
        match fields.find(|&s| s == "-") {
            Some(_) => {}
            None => return None,
        };

        // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - <cgroup> cgroup rw,cpu,cpuacct
        let version = match fields.next() {
            Some("cgroup") => CgroupVersion::V1,
            Some("cgroup2") => CgroupVersion::V2,
            _ => return None,
        };

        // cgroups2 only has a single mount point
        if version == CgroupVersion::V1 {
            // 7 5 0:6 / /sys/fs/cgroup/cpu,cpuacct rw,nosuid,nodev,noexec,relatime shared:7 - cgroup cgroup <rw,cpu,cpuacct>
            let super_opts = some!(fields.nth(1));

            // We only care about the 'cpu' option
            if !super_opts.split(',').any(|opt| opt == "cpu") {
                return None;
            }
        }

        Some(MountInfo {
            version: version,
            root: mnt_root.to_owned(),
            mount_point: mnt_point.to_owned(),
        })
    }
}

impl Subsys {
    fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<Subsys> {
        let file = some!(File::open(proc_path).ok());
        let file = BufReader::new(file);

        file.lines()
            .filter_map(|result| result.ok())
            .filter_map(Subsys::parse_line)
            .fold(None, |previous, line| {
                // already-found v1 trumps v2 since it explicitly specifies its controllers
                if previous.is_some() && line.version == CgroupVersion::V2 {
                    return previous;
                }

                Some(line)
            })
    }

    fn parse_line(line: String) -> Option<Subsys> {
        // Example format:
        // 11:cpu,cpuacct:/
        let mut fields = line.split(':');

        let sub_systems = some!(fields.nth(1));

        let version = if sub_systems.is_empty() {
            CgroupVersion::V2
        } else {
            CgroupVersion::V1
        };

        if version == CgroupVersion::V1 && !sub_systems.split(',').any(|sub| sub == "cpu") {
            return None;
        }

        fields.next().map(|path| Subsys {
            version: version,
            base: path.to_owned(),
        })
    }
}

#[cfg(test)]
mod tests {
    mod v1 {
        use super::super::{Cgroup, CgroupVersion, MountInfo, Subsys};
        use std::path::{Path, PathBuf};

        // `static_in_const` feature is not stable in Rust 1.13.
        static FIXTURES_PROC: &'static str = "fixtures/cgroups/proc/cgroups";

        static FIXTURES_CGROUPS: &'static str = "fixtures/cgroups/cgroups";

        macro_rules! join {
            ($base:expr, $($path:expr),+) => ({
                Path::new($base)
                    $(.join($path))+
            })
        }

        #[test]
        fn test_load_mountinfo() {
            // test only one optional fields
            let path = join!(FIXTURES_PROC, "mountinfo");

            let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V1).unwrap();

            assert_eq!(mnt_info.root, "/");
            assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");

            // test zero optional field
            let path = join!(FIXTURES_PROC, "mountinfo_zero_opt");

            let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V1).unwrap();

            assert_eq!(mnt_info.root, "/");
            assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");

            // test multi optional fields
            let path = join!(FIXTURES_PROC, "mountinfo_multi_opt");

            let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V1).unwrap();

            assert_eq!(mnt_info.root, "/");
            assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup/cpu,cpuacct");
        }

        #[test]
        fn test_load_subsys() {
            let path = join!(FIXTURES_PROC, "cgroup");

            let subsys = Subsys::load_cpu(path).unwrap();

            assert_eq!(subsys.base, "/");
            assert_eq!(subsys.version, CgroupVersion::V1);
        }

        #[test]
        fn test_cgroup_mount() {
            let cases = &[
                ("/", "/sys/fs/cgroup/cpu", "/", Some("/sys/fs/cgroup/cpu")),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd",
                    Some("/sys/fs/cgroup/cpu"),
                ),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd/",
                    Some("/sys/fs/cgroup/cpu"),
                ),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd/large",
                    Some("/sys/fs/cgroup/cpu/large"),
                ),
                // fails
                ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/", None),
                ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/docker", None),
                ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/elsewhere", None),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd-other-dir",
                    None,
                ),
            ];

            for &(root, mount_point, subsys, expected) in cases.iter() {
                let mnt_info = MountInfo {
                    version: CgroupVersion::V1,
                    root: root.into(),
                    mount_point: mount_point.into(),
                };
                let subsys = Subsys {
                    version: CgroupVersion::V1,
                    base: subsys.into(),
                };

                let actual = Cgroup::translate(mnt_info, subsys).map(|c| c.base);
                let expected = expected.map(PathBuf::from);
                assert_eq!(actual, expected);
            }
        }

        #[test]
        fn test_cgroup_cpu_quota() {
            let cgroup = Cgroup::new(CgroupVersion::V1, join!(FIXTURES_CGROUPS, "good"));
            assert_eq!(cgroup.cpu_quota(), Some(6));
        }

        #[test]
        fn test_cgroup_cpu_quota_divide_by_zero() {
            let cgroup = Cgroup::new(CgroupVersion::V1, join!(FIXTURES_CGROUPS, "zero-period"));
            assert!(cgroup.quota_us().is_some());
            assert_eq!(cgroup.period_us(), Some(0));
            assert_eq!(cgroup.cpu_quota(), None);
        }

        #[test]
        fn test_cgroup_cpu_quota_ceil() {
            let cgroup = Cgroup::new(CgroupVersion::V1, join!(FIXTURES_CGROUPS, "ceil"));
            assert_eq!(cgroup.cpu_quota(), Some(2));
        }
    }

    mod v2 {
        use super::super::{Cgroup, CgroupVersion, MountInfo, Subsys};
        use std::path::{Path, PathBuf};

        // `static_in_const` feature is not stable in Rust 1.13.
        static FIXTURES_PROC: &'static str = "fixtures/cgroups2/proc/cgroups";

        static FIXTURES_CGROUPS: &'static str = "fixtures/cgroups2/cgroups";

        macro_rules! join {
            ($base:expr, $($path:expr),+) => ({
                Path::new($base)
                    $(.join($path))+
            })
        }

        #[test]
        fn test_load_mountinfo() {
            // test only one optional fields
            let path = join!(FIXTURES_PROC, "mountinfo");

            let mnt_info = MountInfo::load_cpu(path, CgroupVersion::V2).unwrap();

            assert_eq!(mnt_info.root, "/");
            assert_eq!(mnt_info.mount_point, "/sys/fs/cgroup");
        }

        #[test]
        fn test_load_subsys() {
            let path = join!(FIXTURES_PROC, "cgroup");

            let subsys = Subsys::load_cpu(path).unwrap();

            assert_eq!(subsys.base, "/");
            assert_eq!(subsys.version, CgroupVersion::V2);
        }

        #[test]
        fn test_load_subsys_multi() {
            let path = join!(FIXTURES_PROC, "cgroup_multi");

            let subsys = Subsys::load_cpu(path).unwrap();

            assert_eq!(subsys.base, "/");
            assert_eq!(subsys.version, CgroupVersion::V1);
        }

        #[test]
        fn test_cgroup_mount() {
            let cases = &[
                ("/", "/sys/fs/cgroup/cpu", "/", Some("/sys/fs/cgroup/cpu")),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd",
                    Some("/sys/fs/cgroup/cpu"),
                ),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd/",
                    Some("/sys/fs/cgroup/cpu"),
                ),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd/large",
                    Some("/sys/fs/cgroup/cpu/large"),
                ),
                // fails
                ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/", None),
                ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/docker", None),
                ("/docker/01abcd", "/sys/fs/cgroup/cpu", "/elsewhere", None),
                (
                    "/docker/01abcd",
                    "/sys/fs/cgroup/cpu",
                    "/docker/01abcd-other-dir",
                    None,
                ),
            ];

            for &(root, mount_point, subsys, expected) in cases.iter() {
                let mnt_info = MountInfo {
                    version: CgroupVersion::V1,
                    root: root.into(),
                    mount_point: mount_point.into(),
                };
                let subsys = Subsys {
                    version: CgroupVersion::V1,
                    base: subsys.into(),
                };

                let actual = Cgroup::translate(mnt_info, subsys).map(|c| c.base);
                let expected = expected.map(PathBuf::from);
                assert_eq!(actual, expected);
            }
        }

        #[test]
        fn test_cgroup_cpu_quota() {
            let cgroup = Cgroup::new(CgroupVersion::V2, join!(FIXTURES_CGROUPS, "good"));
            assert_eq!(cgroup.cpu_quota(), Some(6));
        }

        #[test]
        fn test_cgroup_cpu_quota_divide_by_zero() {
            let cgroup = Cgroup::new(CgroupVersion::V2, join!(FIXTURES_CGROUPS, "zero-period"));
            let period = cgroup.max().map(|max| max.1);

            assert_eq!(period, Some(0));
            assert_eq!(cgroup.cpu_quota(), None);
        }

        #[test]
        fn test_cgroup_cpu_quota_ceil() {
            let cgroup = Cgroup::new(CgroupVersion::V2, join!(FIXTURES_CGROUPS, "ceil"));
            assert_eq!(cgroup.cpu_quota(), Some(2));
        }
    }
}

Coverage Report

Created: 2026-07-13 08:11

Line	Count	Source
1		use std::collections::HashMap;
2		use std::fs::File;
3		use std::io::{BufRead, BufReader, Read};
4		use std::mem;
5		use std::path::{Path, PathBuf};
6		use std::sync::atomic::{AtomicUsize, Ordering};
7		use std::sync::Once;
8
9		use libc;
10
11		macro_rules! debug {
12		($($args:expr),*) => ({
13		if false {
14		//if true {
15		println!($($args),*);
16		}
17		});
18		}
19
20		macro_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
32	52.5k	pub fn get_num_cpus() -> usize {
33	52.5k	match cgroups_num_cpus() {
34	0	Some(n) => n,
35	52.5k	None => logical_cpus(),
36		}
37	52.5k	}
38
39	52.5k	fn logical_cpus() -> usize {
40	52.5k	let mut set: libc::cpu_set_t = unsafe { mem::zeroed() };
41	52.5k	if unsafe { libc::sched_getaffinity(0, mem::size_of::<libc::cpu_set_t>(), &mut set) } == 0 {
42	52.5k	let mut count: u32 = 0;
43	53.8M	for i in 0..libc::CPU_SETSIZE as usize {
44	53.8M	if unsafe { libc::CPU_ISSET(i, &set) } {
45	1.68M	count += 1
46	52.1M	}
47		}
48	52.5k	count as usize
49		} else {
50	0	let cpus = unsafe { libc::sysconf(libc::_SC_NPROCESSORS_ONLN) };
51	0	if cpus < 1 {
52	0	1
53		} else {
54	0	cpus as usize
55		}
56		}
57	52.5k	}
58
59	0	pub fn get_num_physical_cpus() -> usize {
60	0	let file = match File::open("/proc/cpuinfo") {
61	0	Ok(val) => val,
62	0	Err(_) => return get_num_cpus(),
63		};
64	0	let reader = BufReader::new(file);
65	0	let mut map = HashMap::new();
66	0	let mut physid: u32 = 0;
67	0	let mut cores: usize = 0;
68	0	let mut chgcount = 0;
69	0	for line in reader.lines().filter_map(\|result\| result.ok()) {
70	0	let mut it = line.split(':');
71	0	let (key, value) = match (it.next(), it.next()) {
72	0	(Some(key), Some(value)) => (key.trim(), value.trim()),
73	0	_ => continue,
74		};
75	0	if key == "physical id" {
76	0	match value.parse() {
77	0	Ok(val) => physid = val,
78	0	Err(_) => break,
79		};
80	0	chgcount += 1;
81	0	}
82	0	if key == "cpu cores" {
83	0	match value.parse() {
84	0	Ok(val) => cores = val,
85	0	Err(_) => break,
86		};
87	0	chgcount += 1;
88	0	}
89	0	if chgcount == 2 {
90	0	map.insert(physid, cores);
91	0	chgcount = 0;
92	0	}
93		}
94	0	let count = map.into_iter().fold(0, \|acc, (_, cores)\| acc + cores);
95
96	0	if count == 0 {
97	0	get_num_cpus()
98		} else {
99	0	count
100		}
101	0	}
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)]
108		static CGROUPS_CPUS: AtomicUsize = ::std::sync::atomic::ATOMIC_USIZE_INIT;
109
110	52.5k	fn cgroups_num_cpus() -> Option<usize> {
111		#[allow(warnings)]
112		static ONCE: Once = ::std::sync::ONCE_INIT;
113
114	52.5k	ONCE.call_once(init_cgroups);
115
116	52.5k	let cpus = CGROUPS_CPUS.load(Ordering::Acquire);
117
118	52.5k	if cpus > 0 {
119	0	Some(cpus)
120		} else {
121	52.5k	None
122		}
123	52.5k	}
124
125	2	fn init_cgroups() {
126		// Should only be called once
127	2	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	2	if cfg!(miri) {
131	0	return;
132	2	}
133
134	2	if let Some(quota) = load_cgroups("/proc/self/cgroup", "/proc/self/mountinfo") {
135	0	if quota == 0 {
136	0	return;
137	0	}
138
139	0	let logical = logical_cpus();
140	0	let count = ::std::cmp::min(quota, logical);
141
142	0	CGROUPS_CPUS.store(count, Ordering::SeqCst);
143	2	}
144	2	}
145
146	2	fn load_cgroups<P1, P2>(cgroup_proc: P1, mountinfo_proc: P2) -> Option<usize>
147	2	where
148	2	P1: AsRef<Path>,
149	2	P2: AsRef<Path>,
150		{
151	2	let subsys = some!(Subsys::load_cpu(cgroup_proc));
152	2	let mntinfo = some!(MountInfo::load_cpu(mountinfo_proc, subsys.version));
153	2	let cgroup = some!(Cgroup::translate(mntinfo, subsys));
154	2	cgroup.cpu_quota()
155	2	}
156
157		#[derive(Debug, Clone, Copy, PartialEq, Eq)]
158		enum CgroupVersion {
159		V1,
160		V2,
161		}
162
163		struct Cgroup {
164		version: CgroupVersion,
165		base: PathBuf,
166		}
167
168		struct MountInfo {
169		version: CgroupVersion,
170		root: String,
171		mount_point: String,
172		}
173
174		struct Subsys {
175		version: CgroupVersion,
176		base: String,
177		}
178
179		impl Cgroup {
180	2	fn new(version: CgroupVersion, dir: PathBuf) -> Cgroup {
181	2	Cgroup { version: version, base: dir }
182	2	}
183
184	2	fn translate(mntinfo: MountInfo, subsys: Subsys) -> Option<Cgroup> {
185		// Translate the subsystem directory via the host paths.
186	2	debug!(
187		"subsys = {:?}; root = {:?}; mount_point = {:?}",
188		subsys.base, mntinfo.root, mntinfo.mount_point
189		);
190
191	2	let rel_from_root = some!(Path::new(&subsys.base).strip_prefix(&mntinfo.root).ok());
192
193	2	debug!("rel_from_root: {:?}", rel_from_root);
194
195		// join(mp.MountPoint, relPath)
196	2	let mut path = PathBuf::from(mntinfo.mount_point);
197	2	path.push(rel_from_root);
198	2	Some(Cgroup::new(mntinfo.version, path))
199	2	}
200
201	2	fn cpu_quota(&self) -> Option<usize> {
202	2	let (quota_us, period_us) = match self.version {
203	2	CgroupVersion::V1 => (some!(self.quota_us()), some!(self.period_us())),
204	0	CgroupVersion::V2 => some!(self.max()),
205		};
206
207		// protect against dividing by zero
208	0	if period_us == 0 {
209	0	return None;
210	0	}
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	0	Some((quota_us as f64 / period_us as f64).ceil() as usize)
216	2	}
217
218	2	fn quota_us(&self) -> Option<usize> {
219	2	self.param("cpu.cfs_quota_us")
220	2	}
221
222	0	fn period_us(&self) -> Option<usize> {
223	0	self.param("cpu.cfs_period_us")
224	0	}
225
226	0	fn max(&self) -> Option<(usize, usize)> {
227	0	let max = some!(self.raw_param("cpu.max"));
228	0	let mut max = some!(max.lines().next()).split(' ');
229
230	0	let quota = some!(max.next().and_then(\|quota\| quota.parse().ok()));
231	0	let period = some!(max.next().and_then(\|period\| period.parse().ok()));
232
233	0	Some((quota, period))
234	0	}
235
236	2	fn param(&self, param: &str) -> Option<usize> {
237	2	let buf = some!(self.raw_param(param));
238
239	2	buf.trim().parse().ok()
240	2	}
241
242	2	fn raw_param(&self, param: &str) -> Option<String> {
243	2	let mut file = some!(File::open(self.base.join(param)).ok());
244
245	2	let mut buf = String::new();
246	2	some!(file.read_to_string(&mut buf).ok());
247
248	2	Some(buf)
249	2	}
250		}
251
252		impl MountInfo {
253	2	fn load_cpu<P: AsRef<Path>>(proc_path: P, version: CgroupVersion) -> Option<MountInfo> {
254	2	let file = some!(File::open(proc_path).ok());
255	2	let file = BufReader::new(file);
256
257	2	file.lines()
258	16	.filter_map(\|result\| result.ok())
259	2	.filter_map(MountInfo::parse_line)
260	2	.find(\|mount_info\| mount_info.version == version)
261	2	}
262
263	16	fn parse_line(line: String) -> Option<MountInfo> {
264	16	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	16	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	16	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	36	match fields.find(\|&s\| s == "-") {
277	16	Some(_) => {}
278	0	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	16	let version = match fields.next() {
283	16	Some("cgroup") => CgroupVersion::V1,
284	12	Some("cgroup2") => CgroupVersion::V2,
285	12	_ => return None,
286		};
287
288		// cgroups2 only has a single mount point
289	4	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	4	let super_opts = some!(fields.nth(1));
292
293		// We only care about the 'cpu' option
294	10	if !super_opts.split(',').any(\|opt\| opt == "cpu") {
295	2	return None;
296	2	}
297	0	}
298
299	2	Some(MountInfo {
300	2	version: version,
301	2	root: mnt_root.to_owned(),
302	2	mount_point: mnt_point.to_owned(),
303	2	})
304	16	}
305		}
306
307		impl Subsys {
308	2	fn load_cpu<P: AsRef<Path>>(proc_path: P) -> Option<Subsys> {
309	2	let file = some!(File::open(proc_path).ok());
310	2	let file = BufReader::new(file);
311
312	2	file.lines()
313	26	.filter_map(\|result\| result.ok())
314	2	.filter_map(Subsys::parse_line)
315	4	.fold(None, \|previous, line\| {
316		// already-found v1 trumps v2 since it explicitly specifies its controllers
317	4	if previous.is_some() && line.version == CgroupVersion::V2 {
318	2	return previous;
319	2	}
320
321	2	Some(line)
322	4	})
323	2	}
324
325	26	fn parse_line(line: String) -> Option<Subsys> {
326		// Example format:
327		// 11:cpu,cpuacct:/
328	26	let mut fields = line.split(':');
329
330	26	let sub_systems = some!(fields.nth(1));
331
332	26	let version = if sub_systems.is_empty() {
333	2	CgroupVersion::V2
334		} else {
335	24	CgroupVersion::V1
336		};
337
338	26	if version == CgroupVersion::V1 && !sub_systems.split(',').any(\|sub\| sub == "cpu") {
339	22	return None;
340	4	}
341
342	4	fields.next().map(\|path\| Subsys {
343	4	version: version,
344	4	base: path.to_owned(),
345	4	})
346	26	}
347		}
348
349		#[cfg(test)]
350		mod 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		}