Preliminary macOS support

README.md

@@ -37,6 +37,28 @@ Complete example [here](https://github.com/elast0ny/affinity-rs/blob/master/exam
 Currently only tested on :
 - Windows
 - Linux (Arch x64)
+- macOS (see note below)
+
+## macOS Caveats
+
+macOS doesn't allow setting thread or process affinities in the same way as Linux and Windows.
+The set_thread_affinity(&cores) call on macOS will only take a single value,
+and it is not a core number, but rather a macOS-specific "affinity tag".
+The following text may be helpful.
+
+From the file <Kernel/thread_policy.h>
+
+> This policy is experimental.
+> 
+> This may be used to express affinity relationships between threads in
+> the task. Threads with the same affinity tag will be scheduled to
+> share an L2 cache if possible. That is, affinity tags are a hint to
+> the scheduler for thread placement.
+> 
+> The namespace of affinity tags is generally local to one task.
+> However, a child task created after the assignment of affinity tags by
+> its parent will share that namespace. In particular, a family of
+> forked processes may be created with a shared affinity namespace.
 
 
 ## License

examples/main.rs

@@ -13,6 +13,7 @@ fn bind_process() -> Result<(), Box<dyn Error>> {
     Ok(())
 }
 
+#[cfg(not(target_os = "macos"))]
 pub fn main() -> Result<(), Box<dyn Error>> {
     println!("Total cores : {}", get_core_num());
 
@@ -31,3 +32,23 @@ pub fn main() -> Result<(), Box<dyn Error>> {
 
     Ok(())
 }
+
+#[cfg(target_os = "macos")]
+pub fn main() -> Result<(), Box<dyn Error>> {
+    println!("Total cores : {}", get_core_num());
+
+    let cores = vec![42]; //42 is an affinity tag, see the readme
+    println!("Binding thread to cores : {:?}", &cores);
+    set_thread_affinity(&cores)?;
+
+    let bound_cores = get_thread_affinity()?;
+    println!("\tCurrent thread affinity : {:?}", bound_cores);
+    println!("\tTotal cores : {}", get_core_num());
+
+    assert_eq!(bound_cores, cores.as_slice());
+
+    #[cfg(target_os = "windows")]
+    bind_process()?;
+
+    Ok(())
+}

src/lib.rs

@@ -39,6 +39,9 @@ cfg_if::cfg_if! {
     } else if #[cfg(target_os = "linux")] {
         mod linux;
         use linux as os;
+    } else if #[cfg(target_os = "macos")] {
+        mod macos;
+        use macos as os;
     } else {
         unimplemented!("This crate does not support your OS yet !");
     }

src/macos.rs

@@ -0,0 +1,84 @@
+use crate::Result;
+use libc::*;
+
+pub fn set_thread_affinity(core_ids: &[usize]) -> Result<()> {
+    if core_ids.len() != 1 {
+        return Err(From::from(
+            "Can only accept a single value/affinity_tag for set_thread_affinity on macos",
+        ));
+    }
+
+    if let Err(e) = _sched_setaffinity(core_ids[0] as i32) {
+        return Err(From::from(format!(
+            "sched_setaffinity failed with errno {}",
+            e
+        )));
+    }
+    Ok(())
+}
+
+pub fn get_thread_affinity() -> Result<Vec<usize>> {
+    let mut affinity = Vec::new();
+    // let mut set: cpu_set_t = unsafe { zeroed() };
+
+    let x = _sched_getaffinity();
+    if let Err(e) = x {
+        return Err(From::from(format!(
+            "sched_getaffinity failed with errno {}",
+            e
+        )));
+    }
+
+    affinity.push(x.unwrap() as usize);
+
+    Ok(affinity)
+}
+
+/* Wrappers around unsafe OS calls */
+fn _sched_setaffinity(affinity_tag: i32) -> std::result::Result<(), i32> {
+    let mut policy_data = thread_affinity_policy_data_t {
+        affinity_tag: affinity_tag,
+    };
+
+    let tid = unsafe { mach_thread_self() };
+
+    let res = unsafe {
+        thread_policy_set(
+            tid,
+            THREAD_AFFINITY_POLICY as u32,
+            (&mut policy_data) as *mut _ as thread_policy_t,
+            1,
+        )
+    };
+    if res != 0 {
+        return Err(errno::errno().into());
+    }
+    Ok(())
+}
+
+fn _sched_getaffinity() -> std::result::Result<i32, i32> {
+    let mut policy_data = thread_affinity_policy_data_t { affinity_tag: -1 };
+
+    let tid = unsafe { mach_thread_self() };
+
+    // false: we want to get the current value, not the default value. If this is `false` after
+    // returning, it means there are no current settings because of other factor, and the
+    // default was returned instead.
+    let mut get_default: boolean_t = 0;
+
+    let mut count: mach_msg_type_number_t = 1;
+    let res = unsafe {
+        thread_policy_get(
+            tid,
+            THREAD_AFFINITY_POLICY as u32,
+            (&mut policy_data) as *mut _ as thread_policy_t,
+            &mut count,
+            &mut get_default,
+        )
+    };
+    if res != 0 {
+        return Err(errno::errno().into());
+    }
+
+    Ok(policy_data.affinity_tag)
+}