use thiserror::Error;
#[derive(Debug, Error)]
pub enum CpuAffinityError {
#[error("cpu affinity is not supported on this platform (Linux only)")]
Unsupported,
#[error("CPU {cpu} out of range (max: {max})")]
CpuOutOfRange {
cpu: usize,
max: usize,
},
#[error("NUMA node {node} not found")]
NumaNodeNotFound {
node: u32,
},
#[error("worker {worker_id} not found in affinity configs (count: {config_count})")]
WorkerNotFound {
worker_id: u32,
config_count: usize,
},
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct CpuId {
pub id: usize,
pub node: u32,
}
impl CpuId {
pub fn new(id: usize, node: u32) -> Self {
Self { id, node }
}
}
pub fn set_thread_affinity(_cpu: usize) -> Result<(), CpuAffinityError> {
Err(CpuAffinityError::Unsupported)
}
pub fn set_thread_affinity_range(_cpus: &[usize]) -> Result<(), CpuAffinityError> {
Err(CpuAffinityError::Unsupported)
}
pub fn get_online_cpus() -> usize {
std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1)
}
#[derive(Debug, Clone)]
pub struct CpuAffinityConfig {
pub worker_id: u32,
pub cpu_ids: Vec<usize>,
pub numa_node: u32,
}
impl CpuAffinityConfig {
pub fn new(worker_id: u32, cpu_ids: Vec<usize>, numa_node: u32) -> Self {
Self { worker_id, cpu_ids, numa_node }
}
pub fn apply(&self) -> Result<(), CpuAffinityError> {
let online = get_online_cpus();
for &cpu in &self.cpu_ids {
if cpu >= online {
return Err(CpuAffinityError::CpuOutOfRange { cpu, max: online });
}
}
Err(CpuAffinityError::Unsupported)
}
}
#[derive(Debug)]
pub struct AffinityManager {
configs: Vec<CpuAffinityConfig>,
}
impl AffinityManager {
pub fn new() -> Self {
Self { configs: Vec::new() }
}
pub fn add_config(&mut self, config: CpuAffinityConfig) {
self.configs.push(config);
}
pub fn apply_all(&self) -> Vec<(u32, Result<(), CpuAffinityError>)> {
self.configs
.iter()
.map(|c| {
let result = c.apply();
(c.worker_id, result)
})
.collect()
}
pub fn apply_for_worker(&self, worker_id: u32) -> Result<(), CpuAffinityError> {
let config = self
.configs
.iter()
.find(|c| c.worker_id == worker_id)
.ok_or(CpuAffinityError::WorkerNotFound {
worker_id,
config_count: self.configs.len(),
})?;
config.apply()
}
pub fn configs(&self) -> &[CpuAffinityConfig] {
&self.configs
}
}
impl Default for AffinityManager {
fn default() -> Self {
Self::new()
}
}