Merge #1148

1148: Implement sched::sched_getaffinity() r=asomers a=thib-ack

Hello,

I found the function `sched::sched_setaffinity()` but I also needed to get the process affinity and I did not find the function `sched::sched_getaffinity()`
So I added the function which maps [sched_getaffinity(2)](https://linux.die.net/man/2/sched_getaffinity) which "get a process's CPU affinity mask" to the sched.rs file.
I hope the code match your guidelines (returned `Result<CpuSet>` instead of pointer parameter)
If that's not the case, tell me, I will fix asap.

I hope this will help !

Thanks,
Thibaut

Co-authored-by: Thibaut Ackermann <thibaut@ackermann.dev>

4 files changed, 114 insertions, 3 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d0f6b087..e4fe67da 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -30,6 +30,9 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 - Added `linkat`
   ([#1101](https://github.com/nix-rust/nix/pull/1101))
 
+- Added `sched_getaffinity`.
+  ([#1148](https://github.com/nix-rust/nix/pull/1148))
+
 ### Changed
 - `sys::socket::recvfrom` now returns
   `Result<(usize, Option<SockAddr>)>` instead of `Result<(usize, SockAddr)>`.
diff --git a/src/sched.rs b/src/sched.rs
index 67188c57..7675dbc2 100644
--- a/src/sched.rs
+++ b/src/sched.rs
@@ -46,6 +46,11 @@ mod sched_linux_like {
 
     pub type CloneCb<'a> = Box<dyn FnMut() -> 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,70 @@ 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<bool> {
-            if field >= 8 * mem::size_of::<libc::cpu_set_t>() {
+            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::<libc::cpu_set_t>() {
+            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::<libc::cpu_set_t>() {
+            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
+        pub fn count() -> usize {
+            8 * mem::size_of::<libc::cpu_set_t>()
+        }
     }
 
+    /// `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,no_run
+    /// 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 +134,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,no_run
+    /// 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<CpuSet> {
+        let mut cpuset = CpuSet::new();
+        let res = unsafe {
+            libc::sched_getaffinity(
+                pid.into(),
+                mem::size_of::<CpuSet>() 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 24260500..4b13fa44 100644
--- a/test/test.rs
+++ b/test/test.rs
@@ -119,6 +119,9 @@ 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",
           target_os = "linux",
diff --git a/test/test_sched.rs b/test/test_sched.rs
new file mode 100644
index 00000000..922196a3
--- /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();
+}