fix macos build

anza-xyz · Dec 8, 2024 · da038b3 · da038b3
1 parent a9f0a51
commit da038b3
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 25 deletions.
diff --git a/thread-manager/src/lib.rs b/thread-manager/src/lib.rs
@@ -58,19 +58,9 @@ impl RuntimeManager {
         self.tokio_runtimes.get(n)
     }
     pub fn set_process_affinity(config: &RuntimeManagerConfig) -> anyhow::Result<Vec<usize>> {
-        let chosen_cores_mask: Vec<usize> = {
-            match config.default_core_allocation {
-                CoreAllocation::PinnedCores { min, max } => (min..max).collect(),
-                CoreAllocation::DedicatedCoreSet { min, max } => (min..max).collect(),
-                CoreAllocation::OsDefault => vec![],
-            }
-        };
+        let chosen_cores_mask = config.default_core_allocation.as_core_mask_vector();
 
-        if cfg!(target_os = "linux") {
-            if let Err(e) = affinity::set_thread_affinity(&chosen_cores_mask) {
-                anyhow::bail!(e.to_string())
-            }
-        }
+        crate::policy::set_thread_affinity(&chosen_cores_mask);
         Ok(chosen_cores_mask)
     }
 
@@ -131,13 +121,14 @@ mod tests {
         std::collections::HashMap,
     };
 
+    // Nobody runs Agave on windows, and on Mac we can not set mask affinity without patching external crate
+    #[cfg(target_os = "linux")]
     fn validate_affinity(expect_cores: &[usize], error_msg: &str) {
-        // Nobody runs Agave on windows, and on Mac we can not set mask affinity without patching external crate
-        if cfg!(target_os = "linux") {
-            let aff = affinity::get_thread_affinity().unwrap();
-            assert_eq!(aff, expect_cores, "{}", error_msg);
-        }
+        let aff = affinity::get_thread_affinity().unwrap();
+        assert_eq!(aff, expect_cores, "{}", error_msg);
     }
+    #[cfg(not(target_os = "linux"))]
+    fn validate_affinity(expect_cores: &[usize], error_msg: &str) {}
 
     #[test]
     fn process_affinity() {

diff --git a/thread-manager/src/policy.rs b/thread-manager/src/policy.rs
@@ -25,6 +25,14 @@ impl CoreAllocation {
     }
 }
 
+#[cfg(target_os = "linux")]
+pub fn set_thread_affinity(cores: &[usize]) {
+    affinity::set_thread_affinity(cores).expect("Can not set thread affinity for runtime worker");
+}
+
+#[cfg(not(target_os = "linux"))]
+pub fn set_thread_affinity(_cores: &[usize]) {}
+
 ///Applies policy to the calling thread
 pub fn apply_policy(
     alloc: &CoreAllocation,
@@ -45,19 +53,13 @@ pub fn apply_policy(
             let core = lg
                 .pop()
                 .expect("Not enough cores provided for pinned allocation");
-            if cfg!(target_os = "linux") {
-                affinity::set_thread_affinity([core])
-                    .expect("Can not set thread affinity for runtime worker");
-            }
+            set_thread_affinity(&[core]);
         }
         CoreAllocation::DedicatedCoreSet { min: _, max: _ } => {
             let lg = chosen_cores_mask
                 .lock()
                 .expect("Can not lock core mask mutex");
-            if cfg!(target_os = "linux") {
-                affinity::set_thread_affinity(&(*lg))
-                    .expect("Can not set thread affinity for runtime worker");
-            }
+            set_thread_affinity(&lg);
         }
         CoreAllocation::OsDefault => {}
     }