SYS-OSS/FRET-LibAFL
4
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

bolts: core affinity simplification for freebsd (#2049)

Browse Source 
since 13.1 (EOL since mid-2023) , supports the sched affinity api like linux/dragonflybsd.
This commit is contained in:
David CARLIER 2024-04-14 09:29:32 +01:00 committed by GitHub
parent 0c68ff033f
commit 4ec50d6a10
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 24 additions and 133 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

157
libafl_bolts/src/core_affinity.rs
View File

@ -215,27 +215,46 @@ pub fn parse_core_bind_arg(args: &str) -> Result<Vec<usize>, Error> {
// Linux Section
#[cfg(any(target_os = "android", target_os = "linux", target_os = "dragonfly"))]
#[cfg(any(
target_os = "android",
target_os = "linux",
target_os = "dragonfly",
target_os = "freebsd"
))]
#[inline]
fn get_core_ids_helper() -> Result<Vec<CoreId>, Error> {
linux::get_core_ids()
}
#[cfg(any(target_os = "android", target_os = "linux", target_os = "dragonfly"))]
#[cfg(any(
target_os = "android",
target_os = "linux",
target_os = "dragonfly",
target_os = "freebsd"
))]
#[inline]
fn set_for_current_helper(core_id: CoreId) -> Result<(), Error> {
linux::set_for_current(core_id)
}
#[cfg(any(target_os = "android", target_os = "linux", target_os = "dragonfly"))]
#[cfg(any(
target_os = "android",
target_os = "linux",
target_os = "dragonfly",
target_os = "freebsd"
))]
mod linux {
use alloc::{string::ToString, vec::Vec};
use std::mem;
#[cfg(not(target_os = "freebsd"))]
use libc::cpu_set_t;
#[cfg(target_os = "freebsd")]
use libc::cpuset_t as cpu_set_t;
#[cfg(target_os = "dragonfly")]
use libc::{cpu_set_t, sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET};
use libc::{sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET};
#[cfg(not(target_os = "dragonfly"))]
use libc::{cpu_set_t, sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET, CPU_SETSIZE};
use libc::{sched_getaffinity, sched_setaffinity, CPU_ISSET, CPU_SET, CPU_SETSIZE};
#[cfg(target_os = "dragonfly")]
const CPU_SETSIZE: libc::c_int = 256;
@ -667,134 +686,6 @@ mod apple {
}
}
// FreeBSD Section
#[cfg(target_os = "freebsd")]
#[inline]
fn get_core_ids_helper() -> Result<Vec<CoreId>, Error> {
freebsd::get_core_ids()
}
#[cfg(target_os = "freebsd")]
#[inline]
fn set_for_current_helper(core_id: CoreId) -> Result<(), Error> {
freebsd::set_for_current(core_id)
}
#[cfg(target_os = "freebsd")]
mod freebsd {
use alloc::vec::Vec;
use std::{mem, thread::available_parallelism};
use libc::{cpuset_setaffinity, cpuset_t, CPU_LEVEL_WHICH, CPU_SET, CPU_WHICH_PID};
use super::CoreId;
use crate::Error;
#[allow(trivial_numeric_casts)]
pub fn get_core_ids() -> Result<Vec<CoreId>, Error> {
Ok((0..(usize::from(available_parallelism()?)))
.map(CoreId)
.collect::<Vec<_>>())
}
pub fn set_for_current(core_id: CoreId) -> Result<(), Error> {
// Turn `core_id` into a `libc::cpuset_t` with only
let mut set = new_cpuset();
unsafe { CPU_SET(core_id.0, &mut set) };
// Set the current thread's core affinity.
let result = unsafe {
cpuset_setaffinity(
CPU_LEVEL_WHICH,
CPU_WHICH_PID,
-1, // Defaults to current thread
mem::size_of::<cpuset_t>(),
&set,
)
};
if result < 0 {
Err(Error::unknown("Failed to set_for_current"))
} else {
Ok(())
}
}
#[cfg(test)]
fn get_affinity_mask() -> Result<cpuset_t, Error> {
let mut set = new_cpuset();
// Try to get current core affinity mask.
let result = unsafe {
libc::cpuset_getaffinity(
CPU_LEVEL_WHICH,
CPU_WHICH_PID,
-1, // Defaults to current thread
mem::size_of::<cpuset_t>(),
&mut set,
)
};
if result == 0 {
Ok(set)
} else {
Err(Error::unknown(
"Failed to retrieve affinity using cpuset_getaffinity",
))
}
}
fn new_cpuset() -> cpuset_t {
unsafe { mem::zeroed::<cpuset_t>() }
}
// FIXME: unstable for now on freebsd
#[cfg(test)]
mod tests {
use libc::CPU_ISSET;
use super::*;
#[test]
#[ignore]
fn test_freebsd_get_affinity_mask() {
get_affinity_mask().unwrap();
}
#[test]
#[ignore]
fn test_freebsd_set_for_current() {
let ids = get_core_ids().unwrap();
assert!(!ids.is_empty());
ids[0].set_affinity().unwrap();
// Ensure that the system pinned the current thread
// to the specified core.
let mut core_mask = new_cpuset();
unsafe { CPU_SET(ids[0].0, &mut core_mask) };
let new_mask = get_affinity_mask().unwrap();
let mut is_equal = true;
for i in 0..256 {
let is_set1 = unsafe { CPU_ISSET(i, &core_mask) };
let is_set2 = unsafe { CPU_ISSET(i, &new_mask) };
if is_set1 != is_set2 {
is_equal = false;
}
}
assert!(is_equal);
}
}
}
// NetBSD Section
#[cfg(target_os = "netbsd")]