From 3d6ba0549f9a21f278db65c5814ae8637b82862a Mon Sep 17 00:00:00 2001 From: Thibaut Ackermann Date: Wed, 30 Oct 2019 21:55:27 +0100 Subject: [PATCH] Implement sched::sched_getaffinity() sched_getaffinity(2) get a process's CPU affinity mask --- CHANGELOG.md | 3 ++ src/sched.rs | 81 ++++++++++++++++++++++++++++++++++++++++++++-- test/test.rs | 3 ++ test/test_sched.rs | 32 ++++++++++++++++++ 4 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 test/test_sched.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index b839a5d315..570eb89947 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,6 +24,9 @@ This project adheres to [Semantic Versioning](http://semver.org/). `Group::from_name`, ([#1139](https://github.com/nix-rust/nix/pull/1139)) +- Added `sched_getaffinity`. + ([#1148](https://github.com/nix-rust/nix/pull/1148)) + ### Changed - `sys::socket::recvfrom` now returns `Result<(usize, Option)>` instead of `Result<(usize, SockAddr)>`. diff --git a/src/sched.rs b/src/sched.rs index 67188c57ee..d818892dac 100644 --- a/src/sched.rs +++ b/src/sched.rs @@ -46,6 +46,11 @@ mod sched_linux_like { pub type CloneCb<'a> = Box isize + 'a>; + /// CpuSet represent a bit-mask of CPUs. + /// CpuSets are used by sched_setaffinity and + /// sched_getaffinity for example. + /// + /// This is a wrapper around `libc::cpu_set_t`. #[repr(C)] #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] pub struct CpuSet { @@ -53,37 +58,72 @@ mod sched_linux_like { } impl CpuSet { + /// Create a new and empty CpuSet. pub fn new() -> CpuSet { CpuSet { cpu_set: unsafe { mem::zeroed() }, } } + /// Test to see if a CPU is in the CpuSet. + /// `field` is the CPU id to test pub fn is_set(&self, field: usize) -> Result { - if field >= 8 * mem::size_of::() { + if field >= CpuSet::count() { Err(Error::Sys(Errno::EINVAL)) } else { Ok(unsafe { libc::CPU_ISSET(field, &self.cpu_set) }) } } + /// Add a CPU to CpuSet. + /// `field` is the CPU id to add pub fn set(&mut self, field: usize) -> Result<()> { - if field >= 8 * mem::size_of::() { + if field >= CpuSet::count() { Err(Error::Sys(Errno::EINVAL)) } else { Ok(unsafe { libc::CPU_SET(field, &mut self.cpu_set) }) } } + /// Remove a CPU from CpuSet. + /// `field` is the CPU id to remove pub fn unset(&mut self, field: usize) -> Result<()> { - if field >= 8 * mem::size_of::() { + CpuSet::count(); + if field >= CpuSet::count() { Err(Error::Sys(Errno::EINVAL)) } else { Ok(unsafe { libc::CPU_CLR(field, &mut self.cpu_set) }) } } + + /// Return the maximum number of CPU in CpuSet + #[inline] + pub fn count() -> usize { + 8 * mem::size_of::() + } } + /// `sched_setaffinity` set a thread's CPU affinity mask + /// ([`sched_setaffinity(2)`](http://man7.org/linux/man-pages/man2/sched_setaffinity.2.html)) + /// + /// `pid` is the thread ID to update. + /// If pid is zero, then the calling thread is updated. + /// + /// The `cpuset` argument specifies the set of CPUs on which the thread + /// will be eligible to run. + /// + /// # Example + /// + /// Binding the current thread to CPU 0 can be done as follows: + /// + /// ```rust + /// use nix::sched::{CpuSet, sched_setaffinity}; + /// use nix::unistd::Pid; + /// + /// let mut cpu_set = CpuSet::new(); + /// cpu_set.set(0); + /// sched_setaffinity(Pid::from_raw(0), &cpu_set); + /// ``` pub fn sched_setaffinity(pid: Pid, cpuset: &CpuSet) -> Result<()> { let res = unsafe { libc::sched_setaffinity( @@ -96,6 +136,41 @@ mod sched_linux_like { Errno::result(res).map(drop) } + /// `sched_getaffinity` get a thread's CPU affinity mask + /// ([`sched_getaffinity(2)`](http://man7.org/linux/man-pages/man2/sched_getaffinity.2.html)) + /// + /// `pid` is the thread ID to check. + /// If pid is zero, then the calling thread is checked. + /// + /// Returned `cpuset` is the set of CPUs on which the thread + /// is eligible to run. + /// + /// # Example + /// + /// Checking if the current thread can run on CPU 0 can be done as follows: + /// + /// ```rust + /// use nix::sched::sched_getaffinity; + /// use nix::unistd::Pid; + /// + /// let cpu_set = sched_getaffinity(Pid::from_raw(0)).unwrap(); + /// if cpu_set.is_set(0).unwrap() { + /// println!("Current thread can run on CPU 0"); + /// } + /// ``` + pub fn sched_getaffinity(pid: Pid) -> Result { + let mut cpuset = CpuSet::new(); + let res = unsafe { + libc::sched_getaffinity( + pid.into(), + mem::size_of::() as libc::size_t, + &mut cpuset.cpu_set, + ) + }; + + Errno::result(res).and(Ok(cpuset)) + } + pub fn clone( mut cb: CloneCb, stack: &mut [u8], diff --git a/test/test.rs b/test/test.rs index 242605005b..4b13fa44db 100644 --- a/test/test.rs +++ b/test/test.rs @@ -118,6 +118,9 @@ mod test_net; mod test_nix_path; mod test_poll; mod test_pty; +#[cfg(any(target_os = "android", + target_os = "linux"))] +mod test_sched; #[cfg(any(target_os = "android", target_os = "freebsd", target_os = "ios", diff --git a/test/test_sched.rs b/test/test_sched.rs new file mode 100644 index 0000000000..922196a3db --- /dev/null +++ b/test/test_sched.rs @@ -0,0 +1,32 @@ +use nix::sched::{sched_getaffinity, sched_setaffinity, CpuSet}; +use nix::unistd::Pid; + +#[test] +fn test_sched_affinity() { + // If pid is zero, then the mask of the calling process is returned. + let initial_affinity = sched_getaffinity(Pid::from_raw(0)).unwrap(); + let mut at_least_one_cpu = false; + let mut last_valid_cpu = 0; + for field in 0..CpuSet::count() { + if initial_affinity.is_set(field).unwrap() { + at_least_one_cpu = true; + last_valid_cpu = field; + } + } + assert!(at_least_one_cpu); + + // Now restrict the running CPU + let mut new_affinity = CpuSet::new(); + new_affinity.set(last_valid_cpu).unwrap(); + sched_setaffinity(Pid::from_raw(0), &new_affinity).unwrap(); + + // And now re-check the affinity which should be only the one we set. + let updated_affinity = sched_getaffinity(Pid::from_raw(0)).unwrap(); + for field in 0..CpuSet::count() { + // Should be set only for the CPU we set previously + assert_eq!(updated_affinity.is_set(field).unwrap(), field==last_valid_cpu) + } + + // Finally, reset the initial CPU set + sched_setaffinity(Pid::from_raw(0), &initial_affinity).unwrap(); +}