From 2ad76496c4b6079e3b8f951bf89982db8f5de48f Mon Sep 17 00:00:00 2001 From: R Aadarsh Date: Sun, 21 Jun 2026 14:01:46 +0530 Subject: [PATCH] * Store addresses as pointers * Remove lots of unsafe code by initialising in the closure passed to `call_once` * Remove `NUMBER_OF_DOMAINS` as size can always be inferred from the hashmap's len --- src/acpi/slit.rs | 26 +++++++++------- src/acpi/srat/mod.rs | 17 +++++----- src/acpi/srat/x86.rs | 27 +++++++++------- src/numa.rs | 74 ++++++++++++++++---------------------------- 4 files changed, 64 insertions(+), 80 deletions(-) diff --git a/src/acpi/slit.rs b/src/acpi/slit.rs index d5c45acfbe..8eafad3318 100644 --- a/src/acpi/slit.rs +++ b/src/acpi/slit.rs @@ -1,7 +1,7 @@ use crate::{ acpi::sdt::Sdt, find_one_sdt, - numa::{self, NumaNode, NUMA_NODES, NUMBER_OF_DOMAINS}, + numa::{self, NumaNode, NUMA_NODES}, }; use core::ops::Add; use hashbrown::HashMap; @@ -10,7 +10,7 @@ use hashbrown::HashMap; pub struct Slit { sdt: &'static Sdt, no: u64, - address: usize, + address: *const u8, } impl Slit { @@ -18,28 +18,30 @@ impl Slit { Self { sdt, no: unsafe { *(sdt.data_address() as *const u64) }, - address: sdt.data_address() + 8, + address: (sdt.data_address() + 8) as *const u8, } } - pub fn init(&self) { - let ndom = *NUMBER_OF_DOMAINS.get().unwrap(); - let address = self.address as *const u8; + pub fn init(&self, numa_nodes: &mut HashMap) { + let address = self.address; + let ndom = NUMA_NODES.get().unwrap().len() as u32; for i in 0..ndom { for j in i..ndom { // ignore distances from a domain to itself, since it is always 10 if i != j { - unsafe { - numa::set_distance(i, j, unsafe { *address.add((i + ndom * j) as usize) }); - numa::set_distance(j, i, unsafe { *address.add((i + ndom * j) as usize) }); - } + numa::set_distance(numa_nodes, i, j, unsafe { + *address.add((i + ndom * j) as usize) + }); + numa::set_distance(numa_nodes, j, i, unsafe { + *address.add((i + ndom * j) as usize) + }); } } } } } -pub fn init() { +pub fn init(numa_nodes: &mut HashMap) { let slit = Slit::new(find_one_sdt!("SLIT")); - slit.init(); + slit.init(numa_nodes); } diff --git a/src/acpi/srat/mod.rs b/src/acpi/srat/mod.rs index 6249c4b9b9..cf88935f00 100644 --- a/src/acpi/srat/mod.rs +++ b/src/acpi/srat/mod.rs @@ -1,9 +1,11 @@ //! See +use hashbrown::HashMap; + use crate::{ acpi::{find_sdt, sdt::Sdt, srat}, find_one_sdt, - numa::{NUMA_NODES, NUMBER_OF_DOMAINS}, + numa::{NumaNode, NUMA_NODES}, }; #[cfg(target_arch = "aarch64")] @@ -17,20 +19,19 @@ mod arch; #[repr(C, packed)] pub struct Srat { sdt: &'static Sdt, - entries: usize, + entries: *const u8, } -pub fn init() { +pub fn init(numa_nodes: &mut HashMap) { let srat = Srat::new(find_one_sdt!("SRAT")); - arch::init_srat(&srat); - NUMBER_OF_DOMAINS.call_once(|| NUMA_NODES.get().unwrap().len() as u32); + arch::init_srat(numa_nodes, &srat); } impl Srat { pub fn new(sdt: &'static Sdt) -> Self { Self { sdt, - entries: sdt.data_address() + 12, + entries: (sdt.data_address() + 12) as *const u8, } } } @@ -55,8 +56,8 @@ impl<'a> Iterator for SratIter<'a> { fn next(&mut self) -> Option { while self.i < self.srat.sdt.data_len() as u32 { - let entry = (self.srat.entries + self.i as usize) as *const u8; - let entry_len = unsafe { *((self.srat.entries + self.i as usize + 1) as *const u8) }; + let entry = unsafe { self.srat.entries.add(self.i as usize) }; + let entry_len = unsafe { *self.srat.entries.add(self.i as usize + 1) }; let entry = Some(match unsafe { *entry } { 0 => SratEntry::LegacyProcessorLocalAffinity(unsafe { diff --git a/src/acpi/srat/x86.rs b/src/acpi/srat/x86.rs index be3b11d6a4..c1b73f973b 100644 --- a/src/acpi/srat/x86.rs +++ b/src/acpi/srat/x86.rs @@ -1,8 +1,10 @@ use core::iter; +use hashbrown::HashMap; + use crate::{ acpi::srat::{to_usize, Srat, SratEntry}, - numa::{self, NUMA_NODES}, + numa::{self, NumaNode, NUMA_NODES}, }; #[inline(always)] @@ -13,23 +15,25 @@ fn to_single_int(high: &[u8; 3], low: u8) -> u32 { u32::from_le_bytes(high_and_low) } -pub fn init_srat(srat: &Srat) { +pub fn init_srat(numa_nodes: &mut HashMap, srat: &Srat) { for affinity in srat { match affinity { - SratEntry::LegacyProcessorLocalAffinity(legacy_processor_local_affinity) => unsafe { + SratEntry::LegacyProcessorLocalAffinity(legacy_processor_local_affinity) => { numa::add_cpu( + numa_nodes, legacy_processor_local_affinity.apic_id as u32, to_single_int( &legacy_processor_local_affinity.proximity_domain_high, legacy_processor_local_affinity.proximity_domain_low, ), ) - }, - SratEntry::MemoryAffinity(memory_affinity) => unsafe { + } + SratEntry::MemoryAffinity(memory_affinity) => { if memory_affinity.length_low == 0 { continue; } numa::add_memory( + numa_nodes, memory_affinity.proximity_domain, to_usize( memory_affinity.base_address_low, @@ -37,13 +41,12 @@ pub fn init_srat(srat: &Srat) { ), to_usize(memory_affinity.length_low, memory_affinity.length_high), ); - }, - SratEntry::ProcessorLocalAffinity(processor_local_affinity) => unsafe { - numa::add_cpu( - processor_local_affinity.x2apic_id, - processor_local_affinity.proximity_domain, - ) - }, + } + SratEntry::ProcessorLocalAffinity(processor_local_affinity) => numa::add_cpu( + numa_nodes, + processor_local_affinity.x2apic_id, + processor_local_affinity.proximity_domain, + ), _ => continue, } } diff --git a/src/numa.rs b/src/numa.rs index d0cae5fef3..33bc3b9ff9 100644 --- a/src/numa.rs +++ b/src/numa.rs @@ -8,7 +8,6 @@ use hashbrown::HashMap; use spin::once::Once; pub static NUMA_NODES: Once> = Once::new(); -pub static NUMBER_OF_DOMAINS: Once = Once::new(); #[derive(Debug)] pub struct NumaMemory { @@ -34,39 +33,28 @@ pub struct NumaNode { } pub fn init() { - NUMA_NODES.call_once(|| HashMap::new()); - let mut flag = false; + NUMA_NODES.call_once(|| { + let mut numa_nodes = HashMap::new(); + #[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))] + { + acpi::srat::init(&mut numa_nodes); + acpi::slit::init(&mut numa_nodes); - #[cfg(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64"))] - { - acpi::srat::init(); - acpi::slit::init(); - flag = true; - } - - #[cfg(any(target_arch = "riscv64", target_arch = "aarch64"))] - { - if !flag { - // todo!() - return; + sort_by_distances(&mut numa_nodes); + reorganise(&mut numa_nodes); + shrink(&mut numa_nodes); } - } - unsafe { - sort_by_distances(); - reorganise(); - shrink(); - } + #[cfg(any(target_arch = "riscv64", target_arch = "aarch64"))] + { + // todo!() + } - // From this point onwards, the global static `NUMA_NODES` or any of its elements - // MUST NOT be mutated by the usual unsafe magic that functions in this file use. + numa_nodes + }); } -pub unsafe fn add_cpu(id: u32, node_id: u32) { - let numa_nodes = NUMA_NODES.get().unwrap(); - - let numa_nodes = unsafe { &mut *(&raw const *numa_nodes as *mut HashMap) }; - +pub fn add_cpu(numa_nodes: &mut HashMap, id: u32, node_id: u32) { if let Some(node) = numa_nodes.get_mut(&id) { node.cpus.push(NumaCpu { id }); } else { @@ -83,11 +71,12 @@ pub unsafe fn add_cpu(id: u32, node_id: u32) { } } -pub unsafe fn add_memory(node_id: u32, start: usize, length: usize) { - let numa_nodes = NUMA_NODES.get().unwrap(); - - let numa_nodes = unsafe { &mut *(&raw const *numa_nodes as *mut HashMap) }; - +pub fn add_memory( + numa_nodes: &mut HashMap, + node_id: u32, + start: usize, + length: usize, +) { if let Some(node) = numa_nodes.get_mut(&node_id) { node.memory.push(NumaMemory { start, length }); } else { @@ -104,18 +93,12 @@ pub unsafe fn add_memory(node_id: u32, start: usize, length: usize) { } } -pub unsafe fn set_distance(src: u32, target: u32, distance: u8) { - let nodes = - unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap) }; - +pub fn set_distance(nodes: &mut HashMap, src: u32, target: u32, distance: u8) { let src = nodes.get_mut(&src).unwrap(); src.distances.push((target, distance)); } -unsafe fn shrink() { - let nodes = - unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap) }; - +fn shrink(nodes: &mut HashMap) { nodes.shrink_to_fit(); for (id, node) in nodes { @@ -130,9 +113,7 @@ unsafe fn shrink() { /// CPUs, the memory is put into a node that has a CPU that is nearest to it. /// /// See the comment above the definition of `NumaNode`. -unsafe fn reorganise() { - let nodes = - unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap) }; +fn reorganise(nodes: &mut HashMap) { let ids = nodes.keys().map(|e| *e).collect::>(); for id in ids { @@ -199,10 +180,7 @@ fn put_for_adoption( } } -unsafe fn sort_by_distances() { - let nodes = - unsafe { &mut *(&raw const *(NUMA_NODES.get().unwrap()) as *mut HashMap) }; - +fn sort_by_distances(nodes: &mut HashMap) { for (id, node) in nodes { node.distances.sort_by_key(|(_, e)| *e); }