Rust demonstration of setting thread affinity with pthread_setaffinity_np

01-rust-thread-affinity.md

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.

02-main.rs

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() });
}