From d6fbf81ad2b4ab69f463bbd938f6406317d96ae1 Mon Sep 17 00:00:00 2001
From: Thibaut Ackermann <thibaut@ackermann.dev>
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       | 13 +++++++++++++
 test/test.rs       |  1 +
 test/test_sched.rs | 36 ++++++++++++++++++++++++++++++++++++
 4 files changed, 53 insertions(+)
 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<SockAddr>)>` instead of `Result<(usize, SockAddr)>`.
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<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 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::<libc::cpu_set_t>()) {
+        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::<libc::cpu_set_t>()) {
+        // 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();
+}