use std::collections::{BTreeMap, BTreeSet, HashSet};
use std::fs;
use std::path::Path;
use anyhow::{Context, Result};
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
use raw_cpuid::CpuId;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CpuGeneration {
Zen2,
Zen3,
Unknown,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RssIrqAffinity {
pub irq: u32,
pub cpus: Vec<usize>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TopologyReport {
pub cpu_vendor: String,
pub family_model: String,
pub cpu_generation: CpuGeneration,
pub numa_nodes: Vec<Vec<usize>>,
pub ccx_groups: Vec<Vec<usize>>,
pub rss_irq_map: Vec<RssIrqAffinity>,
pub selected_cores: Vec<usize>,
}
impl TopologyReport {
pub fn to_json(&self) -> String {
let numa_nodes = self
.numa_nodes
.iter()
.map(|node| format!("[{}]", join_usize(node)))
.collect::<Vec<_>>()
.join(",");
let ccx_groups = self
.ccx_groups
.iter()
.map(|ccx| format!("[{}]", join_usize(ccx)))
.collect::<Vec<_>>()
.join(",");
let rss_irq_map = self
.rss_irq_map
.iter()
.map(|entry| {
format!(
"{{\"irq\":{},\"cpus\":[{}]}}",
entry.irq,
join_usize(&entry.cpus)
)
})
.collect::<Vec<_>>()
.join(",");
format!(
"{{\"cpu_vendor\":\"{}\",\"family_model\":\"{}\",\"cpu_generation\":\"{}\",\"numa_nodes\":[{}],\"ccx_groups\":[{}],\"rss_irq_map\":[{}],\"selected_cores\":[{}]}}",
self.cpu_vendor,
self.family_model,
cpu_generation_label(self.cpu_generation),
numa_nodes,
ccx_groups,
rss_irq_map,
join_usize(&self.selected_cores)
)
}
}
pub fn inspect_topology() -> Result<TopologyReport> {
let (cpu_vendor, display_family, display_model) = read_cpuid_identity();
let family_model = format!("0x{display_family:x}:0x{display_model:x}");
let cpu_generation = detect_cpu_generation(&cpu_vendor, display_family);
let numa_nodes = read_numa_nodes().unwrap_or_default();
let ccx_groups = read_ccx_groups().unwrap_or_default();
let rss_irq_map = read_rss_irq_map().unwrap_or_default();
let cpuset = read_process_cpuset().unwrap_or_else(|_| default_cpuset());
let selected_cores = select_cores(&cpuset, &ccx_groups, &rss_irq_map, 64);
Ok(TopologyReport {
cpu_vendor,
family_model,
cpu_generation,
numa_nodes,
ccx_groups,
rss_irq_map,
selected_cores,
})
}
pub fn detect_cpu_generation(cpu_vendor: &str, display_family: u32) -> CpuGeneration {
if cpu_vendor != "AuthenticAMD" {
return CpuGeneration::Unknown;
}
match display_family {
0x17 => CpuGeneration::Zen2,
0x19 => CpuGeneration::Zen3,
_ => CpuGeneration::Unknown,
}
}
pub fn select_cores(
cpuset: &[usize],
ccx_groups: &[Vec<usize>],
rss_irq_map: &[RssIrqAffinity],
max_cores: usize,
) -> Vec<usize> {
let cpuset_set: BTreeSet<usize> = cpuset.iter().copied().collect();
let rss_union: BTreeSet<usize> = rss_irq_map
.iter()
.flat_map(|entry| entry.cpus.iter().copied())
.collect();
let target: BTreeSet<usize> = if rss_union.is_empty() {
cpuset_set
} else {
cpuset_set
.intersection(&rss_union)
.copied()
.collect::<BTreeSet<_>>()
};
if target.is_empty() {
return Vec::new();
}
let mut selected = Vec::<usize>::new();
let mut used = HashSet::<usize>::new();
for group in ccx_groups {
for &cpu in group {
if target.contains(&cpu) && used.insert(cpu) {
selected.push(cpu);
if selected.len() >= max_cores {
return selected;
}
}
}
}
for cpu in target {
if used.insert(cpu) {
selected.push(cpu);
if selected.len() >= max_cores {
break;
}
}
}
selected
}
pub fn parse_cpu_list(value: &str) -> Result<Vec<usize>> {
let mut cpus = BTreeSet::<usize>::new();
for token in value.trim().split(',').filter(|s| !s.trim().is_empty()) {
let token = token.trim();
if let Some((start, end)) = token.split_once('-') {
let start: usize = start
.trim()
.parse()
.with_context(|| format!("invalid cpu list start: {token}"))?;
let end: usize = end
.trim()
.parse()
.with_context(|| format!("invalid cpu list end: {token}"))?;
if start > end {
anyhow::bail!("invalid cpu range {}-{}", start, end);
}
for cpu in start..=end {
cpus.insert(cpu);
}
} else {
let cpu: usize = token
.parse()
.with_context(|| format!("invalid cpu list token: {token}"))?;
cpus.insert(cpu);
}
}
Ok(cpus.into_iter().collect())
}
fn read_numa_nodes() -> Result<Vec<Vec<usize>>> {
let node_root = Path::new("/sys/devices/system/node");
if !node_root.exists() {
return Ok(Vec::new());
}
let mut nodes = Vec::<Vec<usize>>::new();
for entry in fs::read_dir(node_root).context("failed reading numa node root")? {
let entry = entry.context("failed reading numa node entry")?;
let name = entry.file_name();
let name = name.to_string_lossy();
if !name.starts_with("node") {
continue;
}
let cpulist = entry.path().join("cpulist");
let text = match fs::read_to_string(&cpulist) {
Ok(text) => text,
Err(_) => continue,
};
let parsed = parse_cpu_list(text.trim())?;
if !parsed.is_empty() {
nodes.push(parsed);
}
}
nodes.sort();
Ok(nodes)
}
fn read_ccx_groups() -> Result<Vec<Vec<usize>>> {
let cpu_root = Path::new("/sys/devices/system/cpu");
if !cpu_root.exists() {
return Ok(Vec::new());
}
let mut ccx_map = BTreeMap::<String, BTreeSet<usize>>::new();
for entry in fs::read_dir(cpu_root).context("failed reading cpu root")? {
let entry = entry.context("failed reading cpu entry")?;
let file_name = entry.file_name();
let file_name = file_name.to_string_lossy();
if !file_name.starts_with("cpu") {
continue;
}
let cpu_id = match file_name.trim_start_matches("cpu").parse::<usize>() {
Ok(id) => id,
Err(_) => continue,
};
let cache_index3 = entry.path().join("cache/index3");
if !cache_index3.exists() {
continue;
}
let cache_id = fs::read_to_string(cache_index3.join("id"))
.unwrap_or_else(|_| "unknown".to_string())
.trim()
.to_string();
let shared_cpu_list = fs::read_to_string(cache_index3.join("shared_cpu_list"))
.unwrap_or_else(|_| cpu_id.to_string());
let shared = parse_cpu_list(shared_cpu_list.trim()).unwrap_or_else(|_| vec![cpu_id]);
let group = ccx_map.entry(cache_id).or_default();
for cpu in shared {
group.insert(cpu);
}
}
let mut groups = ccx_map
.into_values()
.map(|set| set.into_iter().collect::<Vec<_>>())
.collect::<Vec<_>>();
groups.sort();
Ok(groups)
}
fn read_rss_irq_map() -> Result<Vec<RssIrqAffinity>> {
let interrupts =
fs::read_to_string("/proc/interrupts").context("failed reading /proc/interrupts")?;
let mut map = Vec::<RssIrqAffinity>::new();
for line in interrupts.lines() {
let trimmed = line.trim_start();
let Some((irq_text, rest)) = trimmed.split_once(':') else {
continue;
};
let irq = match irq_text.trim().parse::<u32>() {
Ok(irq) => irq,
Err(_) => continue,
};
let lower = rest.to_ascii_lowercase();
if !(lower.contains("eth")
|| lower.contains("ens")
|| lower.contains("enp")
|| lower.contains("eno")
|| lower.contains("mlx")
|| lower.contains("net"))
{
continue;
}
let affinity_path = format!("/proc/irq/{irq}/smp_affinity_list");
let affinity = match fs::read_to_string(&affinity_path) {
Ok(text) => text,
Err(_) => continue,
};
let cpus = parse_cpu_list(affinity.trim()).unwrap_or_default();
if cpus.is_empty() {
continue;
}
map.push(RssIrqAffinity { irq, cpus });
}
map.sort_by_key(|entry| entry.irq);
Ok(map)
}
fn read_process_cpuset() -> Result<Vec<usize>> {
let status =
fs::read_to_string("/proc/self/status").context("failed reading /proc/self/status")?;
for line in status.lines() {
if let Some(value) = line.strip_prefix("Cpus_allowed_list:\t") {
return parse_cpu_list(value.trim());
}
}
anyhow::bail!("Cpus_allowed_list not found in /proc/self/status")
}
fn default_cpuset() -> Vec<usize> {
let n = std::thread::available_parallelism()
.map(|n| n.get())
.unwrap_or(1);
(0..n).collect()
}
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
fn read_cpuid_identity() -> (String, u32, u32) {
let cpuid = CpuId::new();
let cpu_vendor = cpuid
.get_vendor_info()
.map(|v| v.as_str().to_string())
.unwrap_or_else(|| "unknown".to_string());
let (display_family, display_model) = cpuid
.get_feature_info()
.map(|f| {
let family = f.family_id() as u32;
let model = f.model_id() as u32;
let ext_family = f.extended_family_id() as u32;
let ext_model = f.extended_model_id() as u32;
let resolved_family = if family == 0xF {
family + ext_family
} else {
family
};
let resolved_model = if family == 0x6 || family == 0xF {
(ext_model << 4) + model
} else {
model
};
(resolved_family, resolved_model)
})
.unwrap_or((0, 0));
(cpu_vendor, display_family, display_model)
}
#[cfg(not(any(target_arch = "x86", target_arch = "x86_64")))]
fn read_cpuid_identity() -> (String, u32, u32) {
("unknown".to_string(), 0, 0)
}
fn cpu_generation_label(value: CpuGeneration) -> &'static str {
match value {
CpuGeneration::Zen2 => "Zen2",
CpuGeneration::Zen3 => "Zen3",
CpuGeneration::Unknown => "Unknown",
}
}
fn join_usize(values: &[usize]) -> String {
values
.iter()
.map(|v| v.to_string())
.collect::<Vec<_>>()
.join(",")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_cpu_list_supports_ranges_and_singletons() {
let parsed = parse_cpu_list("0-3,8,10-12").expect("cpu list parse should succeed");
assert_eq!(parsed, vec![0, 1, 2, 3, 8, 10, 11, 12]);
}
#[test]
fn detect_cpu_generation_maps_amd_families() {
assert_eq!(
detect_cpu_generation("AuthenticAMD", 0x17),
CpuGeneration::Zen2
);
assert_eq!(
detect_cpu_generation("AuthenticAMD", 0x19),
CpuGeneration::Zen3
);
assert_eq!(
detect_cpu_generation("AuthenticAMD", 0x1A),
CpuGeneration::Unknown
);
}
#[test]
fn select_cores_prefers_rss_intersection() {
let cpuset = vec![0, 1, 2, 3, 4, 5];
let ccx_groups = vec![vec![0, 1, 2], vec![3, 4, 5]];
let rss = vec![RssIrqAffinity {
irq: 44,
cpus: vec![1, 2, 4],
}];
let selected = select_cores(&cpuset, &ccx_groups, &rss, 64);
assert_eq!(selected, vec![1, 2, 4]);
}
}