1use nix::sched::{sched_getaffinity, sched_getcpu, sched_setaffinity, CpuSet}; 2use nix::unistd::Pid; 3 4#[test] 5fn test_sched_affinity() { 6 // If pid is zero, then the mask of the calling process is returned. 7 let initial_affinity = sched_getaffinity(Pid::from_raw(0)).unwrap(); 8 let mut at_least_one_cpu = false; 9 let mut last_valid_cpu = 0; 10 for field in 0..CpuSet::count() { 11 if initial_affinity.is_set(field).unwrap() { 12 at_least_one_cpu = true; 13 last_valid_cpu = field; 14 } 15 } 16 assert!(at_least_one_cpu); 17 18 // Now restrict the running CPU 19 let mut new_affinity = CpuSet::new(); 20 new_affinity.set(last_valid_cpu).unwrap(); 21 sched_setaffinity(Pid::from_raw(0), &new_affinity).unwrap(); 22 23 // And now re-check the affinity which should be only the one we set. 24 let updated_affinity = sched_getaffinity(Pid::from_raw(0)).unwrap(); 25 for field in 0..CpuSet::count() { 26 // Should be set only for the CPU we set previously 27 assert_eq!( 28 updated_affinity.is_set(field).unwrap(), 29 field == last_valid_cpu 30 ) 31 } 32 33 // Now check that we're also currently running on the CPU in question. 34 let cur_cpu = sched_getcpu().unwrap(); 35 assert_eq!(cur_cpu, last_valid_cpu); 36 37 // Finally, reset the initial CPU set 38 sched_setaffinity(Pid::from_raw(0), &initial_affinity).unwrap(); 39} 40