diff --git a/CHANGELOG.md b/CHANGELOG.md index b839a5d315..5aabaee616 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,8 @@ This project adheres to [Semantic Versioning](http://semver.org/). - Added `User::from_uid`, `User::from_name`, `User::from_gid` and `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 diff --git a/src/sched.rs b/src/sched.rs index 67188c57ee..bcd060d953 100644 --- a/src/sched.rs +++ b/src/sched.rs @@ -96,6 +96,19 @@ mod sched_linux_like { Errno::result(res).map(drop) } + 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..a46f5c2abf 100644 --- a/test/test.rs +++ b/test/test.rs @@ -118,6 +118,7 @@ mod test_net; mod test_nix_path; mod test_poll; mod test_pty; +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..d7c48b6ae7 --- /dev/null +++ b/test/test_sched.rs @@ -0,0 +1,36 @@ + +#[cfg(any(target_os = "android", + target_os = "linux"))] +#[test] +fn test_sched_affinity() { + use nix::sched::{sched_getaffinity, sched_setaffinity, CpuSet}; + use nix::unistd::Pid; + use std::mem; + + // 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..(8 * mem::size_of::()) { + 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..(8 * mem::size_of::()) { + // 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(); +}