diff --git a/rust/scx_utils/src/lib.rs b/rust/scx_utils/src/lib.rs index 5dcc2bf..010f78a 100644 --- a/rust/scx_utils/src/lib.rs +++ b/rust/scx_utils/src/lib.rs @@ -61,6 +61,7 @@ pub use topology::Core; pub use topology::Cpu; pub use topology::Node; pub use topology::Topology; +pub use topology::TopologyMap; mod cpumask; pub use cpumask::Cpumask; diff --git a/rust/scx_utils/src/topology.rs b/rust/scx_utils/src/topology.rs index 09ca1d6..8b68394 100644 --- a/rust/scx_utils/src/topology.rs +++ b/rust/scx_utils/src/topology.rs @@ -75,6 +75,7 @@ use glob::glob; use sscanf::sscanf; use std::collections::BTreeMap; use std::path::Path; +use std::slice::Iter; #[derive(Debug, Clone)] pub struct Cpu { @@ -249,6 +250,53 @@ impl Topology { } } +/// Generate a topology map from a Topology object, represented as an array of arrays. +/// +/// Each inner array corresponds to a core containing its associated CPU IDs. This map can +/// facilitate efficient iteration over the host's topology. +/// +/// # Example +/// +/// ``` +/// let topo = Topology::new()?; +/// let topo_map = TopologyMap::new(topo)?; +/// +/// for (core_id, core) in topo_map.iter().enumerate() { +/// for cpu in core { +/// println!("core={} cpu={}", core_id, cpu); +/// } +/// } +/// ``` +#[derive(Debug)] +pub struct TopologyMap { + map: Vec>, + nr_cpus_possible: usize, +} + +impl TopologyMap { + pub fn new(topo: Topology) -> Result { + let mut map: Vec> = Vec::new(); + + for core in topo.cores().into_iter() { + let mut cpu_ids: Vec = Vec::new(); + for cpu_id in core.span().clone().into_iter() { + cpu_ids.push(cpu_id); + } + map.push(cpu_ids); + } + let nr_cpus_possible = topo.nr_cpus_possible; + + Ok(TopologyMap { map, nr_cpus_possible, }) + } + + pub fn nr_cpus_possible(&self) -> usize { + self.nr_cpus_possible + } + + pub fn iter(&self) -> Iter> { + self.map.iter() + } +} /********************************************** * Helper functions for creating the Topology * diff --git a/scheds/rust/scx_rustland/src/main.rs b/scheds/rust/scx_rustland/src/main.rs index 13b4cad..9a0dd3b 100644 --- a/scheds/rust/scx_rustland/src/main.rs +++ b/scheds/rust/scx_rustland/src/main.rs @@ -10,6 +10,7 @@ mod bpf; use bpf::*; use scx_utils::Topology; +use scx_utils::TopologyMap; use std::thread; @@ -240,7 +241,7 @@ impl TaskTree { // Main scheduler object struct Scheduler<'a> { bpf: BpfScheduler<'a>, // BPF connector - topo: Topology, // Host topology + topo_map: TopologyMap, // Host topology task_pool: TaskTree, // tasks ordered by vruntime task_map: TaskInfoMap, // map pids to the corresponding task information min_vruntime: u64, // Keep track of the minimum vruntime across all tasks @@ -256,6 +257,7 @@ impl<'a> Scheduler<'a> { fn init(opts: &Opts) -> Result { // Initialize core mapping topology. let topo = Topology::new().expect("Failed to build host topology"); + let topo_map = TopologyMap::new(topo).expect("Failed to generate topology map"); // Save the default time slice (in ns) in the scheduler class. let slice_ns = opts.slice_us * NSEC_PER_USEC; @@ -283,7 +285,7 @@ impl<'a> Scheduler<'a> { let init_page_faults: u64 = 0; // Low-level BPF connector. - let nr_online_cpus = topo.span().weight(); + let nr_online_cpus = topo_map.nr_cpus_possible(); let bpf = BpfScheduler::init( opts.slice_us, nr_online_cpus as i32, @@ -297,7 +299,7 @@ impl<'a> Scheduler<'a> { // Return scheduler object. Ok(Self { bpf, - topo, + topo_map, task_pool, task_map, min_vruntime, @@ -318,10 +320,10 @@ impl<'a> Scheduler<'a> { let mut idle_cpu_count = 0; // Count the number of cores where all the CPUs are idle. - for core in self.topo.cores().iter() { + for core in self.topo_map.iter() { let mut all_idle = true; - for cpu_id in core.span().clone().into_iter() { - if self.bpf.get_cpu_pid(cpu_id as i32) != 0 { + for cpu_id in core { + if self.bpf.get_cpu_pid(*cpu_id as i32) != 0 { all_idle = false; break; } @@ -381,7 +383,7 @@ impl<'a> Scheduler<'a> { // // This allows to survive stress tests that are spawning a massive amount of // tasks. - self.eff_slice_boost = (self.slice_boost * self.topo.nr_cpus_possible() as u64 + self.eff_slice_boost = (self.slice_boost * self.topo_map.nr_cpus_possible() as u64 / self.task_pool.tasks.len().max(1) as u64) .max(1); @@ -706,14 +708,14 @@ impl<'a> Scheduler<'a> { Err(_) => -1, }; info!("Running tasks:"); - for core in self.topo.cores().iter() { - for (cpu_id, _) in core.cpus().iter() { + for (core_id, core) in self.topo_map.iter().enumerate() { + for cpu_id in core { let pid = if *cpu_id as i32 == sched_cpu { "[self]".to_string() } else { self.bpf.get_cpu_pid(*cpu_id as i32).to_string() }; - info!(" core {:2} cpu {:2} pid={}", core.id(), cpu_id, pid); + info!(" core {:2} cpu {:2} pid={}", core_id, cpu_id, pid); } }