From 5be1beb1f252b99ce67547863c3f24ea9e9426cf Mon Sep 17 00:00:00 2001
From: kitcatier <127749832+kitcatier@users.noreply.github.com>
Date: Sat, 22 Apr 2023 10:44:06 +0800
Subject: [PATCH] make set_for_current function unsafe

---
 src/bastion-executor/benches/stats.rs      |  2 +-
 src/bastion-executor/src/placement.rs      | 22 +++++++++++-----------
 src/bastion-executor/src/thread_manager.rs |  2 +-
 3 files changed, 13 insertions(+), 13 deletions(-)

diff --git a/src/bastion-executor/benches/stats.rs b/src/bastion-executor/benches/stats.rs
index 684e7cb1..bdb19a2b 100644
--- a/src/bastion-executor/benches/stats.rs
+++ b/src/bastion-executor/benches/stats.rs
@@ -10,7 +10,7 @@ fn stress_stats<S: SmpStats + Sync + Send>(stats: &'static S) {
     let mut handles = Vec::with_capacity(*core_count());
     for core in get_cores() {
         let handle = thread::spawn(move || {
-            placement::set_for_current(*core);
+            unsafe { placement::set_for_current(*core) };
             for i in 0..100 {
                 stats.store_load(core.id, 10);
                 if i % 3 == 0 {
diff --git a/src/bastion-executor/src/placement.rs b/src/bastion-executor/src/placement.rs
index 1d3e0a76..ba1fa1c8 100644
--- a/src/bastion-executor/src/placement.rs
+++ b/src/bastion-executor/src/placement.rs
@@ -17,7 +17,7 @@ pub fn get_num_cores() -> Option<usize> {
 }
 ///
 /// Sets the current threads affinity
-pub fn set_for_current(core_id: CoreId) {
+pub unsafe fn set_for_current(core_id: CoreId) {
     tracing::trace!("Executor: placement: set affinity on core {}", core_id.id);
     set_for_current_helper(core_id);
 }
@@ -41,7 +41,7 @@ fn get_core_ids_helper() -> Option<Vec<CoreId>> {
 #[cfg(target_os = "linux")]
 #[inline]
 fn set_for_current_helper(core_id: CoreId) {
-    linux::set_for_current(core_id);
+    unsafe { linux::set_for_current(core_id) };
 }
 
 #[cfg(target_os = "linux")]
@@ -68,7 +68,7 @@ mod linux {
         }
     }
 
-    pub fn set_for_current(core_id: CoreId) {
+    pub unsafe fn set_for_current(core_id: CoreId) {
         // Turn `core_id` into a `libc::cpu_set_t` with only
         // one core active.
         let mut set = new_cpu_set();
@@ -141,7 +141,7 @@ mod linux {
 
             assert!(!ids.is_empty());
 
-            set_for_current(ids[0]);
+            unsafe { set_for_current(ids[0]) };
 
             // Ensure that the system pinned the current thread
             // to the specified core.
@@ -177,7 +177,7 @@ fn get_core_ids_helper() -> Option<Vec<CoreId>> {
 #[cfg(target_os = "windows")]
 #[inline]
 fn set_for_current_helper(core_id: CoreId) {
-    windows::set_for_current(core_id);
+    unsafe { windows::set_for_current(core_id);
 }
 
 #[cfg(target_os = "windows")]
@@ -211,7 +211,7 @@ mod windows {
         }
     }
 
-    pub fn set_for_current(core_id: CoreId) {
+    pub unsafe fn set_for_current(core_id: CoreId) {
         // Convert `CoreId` back into mask.
         let mask: DWORD_PTR = 1 << core_id.id;
 
@@ -267,7 +267,7 @@ mod windows {
 
             assert!(ids.len() > 0);
 
-            set_for_current(ids[0]);
+            unsafe { set_for_current(ids[0]) };
         }
     }
 }
@@ -283,7 +283,7 @@ fn get_core_ids_helper() -> Option<Vec<CoreId>> {
 #[cfg(target_os = "macos")]
 #[inline]
 fn set_for_current_helper(core_id: CoreId) {
-    macos::set_for_current(core_id);
+    unsafe { macos::set_for_current(core_id) };
 }
 
 #[cfg(target_os = "macos")]
@@ -328,7 +328,7 @@ mod macos {
         )
     }
 
-    pub fn set_for_current(core_id: CoreId) {
+    pub unsafe fn set_for_current(core_id: CoreId) {
         let thread_affinity_policy_count: MachMsgTypeNumberT =
             mem::size_of::<ThreadAffinityPolicyDataT>() as MachMsgTypeNumberT
                 / mem::size_of::<IntegerT>() as MachMsgTypeNumberT;
@@ -370,7 +370,7 @@ mod macos {
 
             assert!(ids.len() > 0);
 
-            set_for_current(ids[0]);
+            unsafe { set_for_current(ids[0]) };
         }
     }
 }
@@ -410,6 +410,6 @@ mod tests {
 
         assert!(!ids.is_empty());
 
-        set_for_current(ids[0]);
+        unsafe { set_for_current(ids[0]) };
     }
 }
diff --git a/src/bastion-executor/src/thread_manager.rs b/src/bastion-executor/src/thread_manager.rs
index 4f89dc08..a6e46727 100644
--- a/src/bastion-executor/src/thread_manager.rs
+++ b/src/bastion-executor/src/thread_manager.rs
@@ -298,7 +298,7 @@ impl DynamicPoolManager {
     fn affinity_pinner() {
         if 1 != *load_balancer::core_count() {
             let mut core = ROUND_ROBIN_PIN.lock().unwrap();
-            placement::set_for_current(*core);
+            unsafe { placement::set_for_current(*core) };
             core.id = (core.id + 1) % *load_balancer::core_count();
         }
     }