Просмотр исходного кода

minerd: Add functions to set thread affinity

x 7 месяцев назад
Родитель
Сommit
21ac62a430
1 измененных файлов с 67 добавлено и 0 удалено
  1. 67 0
      bin/minerd/src/cpu.rs

+ 67 - 0
bin/minerd/src/cpu.rs

@@ -16,6 +16,8 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
+use std::io;
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub struct CpuThread {
     affinity: i32,
@@ -100,3 +102,68 @@ pub fn get_affinity(index: u64, affinity: i32) -> i32 {
 
     -1
 }
+
+/// Binds the current thread to the specified core(s)
+pub fn set_thread_affinity<B: AsRef<[usize]>>(core_ids: B) -> io::Result<()> {
+    os::set_thread_affinity(core_ids.as_ref())
+}
+
+/// Returns a list of cores that the current thread is bound to
+pub fn get_thread_affinity() -> io::Result<Vec<usize>> {
+    os::get_thread_affinity()
+}
+
+// https://github.com/elast0ny/affinity/
+// Licensed under MIT
+#[cfg(target_os = "linux")]
+mod os {
+    use libc::{
+        cpu_set_t, pid_t, sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET, CPU_SETSIZE,
+    };
+    use std::{
+        io,
+        mem::{size_of, zeroed},
+    };
+
+    pub(super) fn set_thread_affinity(core_ids: &[usize]) -> io::Result<()> {
+        let mut set: cpu_set_t = unsafe { zeroed() };
+        unsafe {
+            for core_id in core_ids {
+                CPU_SET(*core_id, &mut set);
+            }
+        }
+
+        _sched_setaffinity(0, size_of::<cpu_set_t>(), &set)
+    }
+
+    pub(super) fn get_thread_affinity() -> io::Result<Vec<usize>> {
+        let mut affinity = vec![];
+        let mut set: cpu_set_t = unsafe { zeroed() };
+
+        _sched_getaffinity(0, size_of::<cpu_set_t>(), &mut set)?;
+
+        for i in 0..CPU_SETSIZE as usize {
+            if unsafe { CPU_ISSET(i, &set) } {
+                affinity.push(i);
+            }
+        }
+
+        Ok(affinity)
+    }
+
+    fn _sched_setaffinity(pid: pid_t, cpusetsize: usize, mask: &cpu_set_t) -> io::Result<()> {
+        let res = unsafe { sched_setaffinity(pid, cpusetsize, mask) };
+        if res != 0 {
+            return Err(io::Error::last_os_error())
+        }
+        Ok(())
+    }
+
+    fn _sched_getaffinity(pid: pid_t, cpusetsize: usize, mask: &mut cpu_set_t) -> io::Result<()> {
+        let res = unsafe { sched_getaffinity(pid, cpusetsize, mask) };
+        if res != 0 {
+            return Err(io::Error::last_os_error())
+        }
+        Ok(())
+    }
+}