Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Rust demonstration of setting thread affinity with pthread_setaffinity_np

Call this program with -s N to set a specific CPU. Without -s, the thread will have the default affinity.

This program will print out the main threads CPU, the child threads CPU and the child threads affinity after optionally setting it.

extern crate libc;
use std::thread;
use std::sync;
fn some_thread(b: sync::Arc<sync::Barrier>) {
let mut i = 0;
// Wait for affinity to be set
b.wait();
while i < 10 {
let cpu = unsafe { libc::sched_getcpu() };
println!("{} (on CPU {})", i, cpu);
i += 1;
thread::sleep(std::time::Duration::from_millis(100));
}
}
fn main() {
use std::os::unix::thread::JoinHandleExt;
let barrier = sync::Arc::new(sync::Barrier::new(2));
let barrier_1 = barrier.clone();
let handle = thread::spawn(move || some_thread(barrier_1));
let pthread = handle.as_pthread_t();
println!("-> Main Thread on CPU {}", unsafe { libc::sched_getcpu() });
thread::sleep(std::time::Duration::from_millis(10));
println!("-> Main Thread on CPU {}", unsafe { libc::sched_getcpu() });
let mut args = std::env::args();
args.next();
if let Some(true) = args.next().map(|v| v == "-s") {
// Set thread affinity before startup
let cpu: usize = args.next().unwrap().parse().unwrap();
println!("Setting affinity for cpu {}", cpu);
unsafe {
let mut set: libc::cpu_set_t = std::mem::zeroed();
libc::CPU_SET(cpu, &mut set);
libc::pthread_setaffinity_np(pthread, std::mem::size_of::<libc::cpu_set_t>(), &set);
}
}
// Signal, that affinity has been set
barrier.wait();
println!("-> Main Thread on CPU {}", unsafe { libc::sched_getcpu() });
{
let num_cpus = libc::CPU_SETSIZE as usize;
println!("CPUS: {}, Struct size: {}", num_cpus, std::mem::size_of::<libc::cpu_set_t>());
let set = unsafe {
let mut set: libc::cpu_set_t = std::mem::zeroed();
libc::pthread_getaffinity_np(pthread, std::mem::size_of::<libc::cpu_set_t>(), &mut set);
set
};
let cpuset = (0..num_cpus)
.filter_map(|cpu| if unsafe { libc::CPU_ISSET(cpu, &set) } { Some(cpu) } else { None })
.collect::<Vec<_>>();
println!("{:?}", cpuset);
}
println!("-> Main Thread on CPU {}", unsafe { libc::sched_getcpu() });
handle.join().unwrap();
println!("-> Main Thread on CPU {}", unsafe { libc::sched_getcpu() });
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment