|
@@ -16,6 +16,8 @@
|
|
|
* along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
* along with this program. If not, see <https://www.gnu.org/licenses/>.
|
|
|
*/
|
|
*/
|
|
|
|
|
|
|
|
|
|
+use std::io;
|
|
|
|
|
+
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
|
pub struct CpuThread {
|
|
pub struct CpuThread {
|
|
|
affinity: i32,
|
|
affinity: i32,
|
|
@@ -100,3 +102,68 @@ pub fn get_affinity(index: u64, affinity: i32) -> i32 {
|
|
|
|
|
|
|
|
-1
|
|
-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(())
|
|
|
|
|
+ }
|
|
|
|
|
+}
|