Refactor the this_cpu API with PinCurrentCpu

ostd/src/cpu/local/cpu_local.rs

@@ -23,7 +23,7 @@ use crate::{
 /// # Example
 ///
 /// ```rust
-/// use ostd::{cpu_local, cpu::this_cpu};
+/// use ostd::{cpu_local, cpu::PinCurrentCpu, task::disable_preempt};
 /// use core::{sync::atomic::{AtomicU32, Ordering}, cell::Cell};
 ///
 /// cpu_local! {
@@ -32,16 +32,12 @@ use crate::{
 /// }
 ///
 /// fn not_an_atomic_function() {
-///     let ref_of_foo = FOO.get_on_cpu(this_cpu());
-///     // Note that the value of `FOO` here doesn't necessarily equal to the value
-///     // of `FOO` of exactly the __current__ CPU. Since that task may be preempted
-///     // and moved to another CPU since `ref_of_foo` is created.
+///     let preempt_guard = disable_preempt();
+///     let ref_of_foo = FOO.get_on_cpu(preempt_guard.current_cpu());
 ///     let val_of_foo = ref_of_foo.load(Ordering::Relaxed);
 ///     println!("FOO VAL: {}", val_of_foo);
 ///
 ///     let bar_guard = BAR.borrow_irq_disabled();
-///     // Here the value of `BAR` is always the one in the __current__ CPU since
-///     // interrupts are disabled and we do not explicitly yield execution here.
 ///     let val_of_bar = bar_guard.get();
 ///     println!("BAR VAL: {}", val_of_bar);
 /// }

ostd/src/cpu/local/mod.rs

@@ -98,7 +98,7 @@ pub unsafe fn init_on_bsp() {
     let num_cpus = super::num_cpus();
 
     let mut cpu_local_storages = Vec::with_capacity(num_cpus as usize - 1);
-    for cpu_i in 1..num_cpus {
+    for _ in 1..num_cpus {
         let ap_pages = {
             let nbytes = (bsp_end_va - bsp_base_va).align_up(PAGE_SIZE);
             page::allocator::alloc_contiguous(nbytes, |_| KernelMeta::default()).unwrap()
@@ -116,23 +116,11 @@ pub unsafe fn init_on_bsp() {
             );
         }
 
-        // SAFETY: bytes `0:4` are reserved for storing CPU ID.
-        unsafe {
-            (ap_pages_ptr as *mut u32).write(cpu_i);
-        }
-
         cpu_local_storages.push(ap_pages);
     }
 
     CPU_LOCAL_STORAGES.call_once(|| cpu_local_storages);
 
-    // Write the CPU ID of BSP to the first 4 bytes of the CPU-local area.
-    let bsp_cpu_id_ptr = bsp_base_va as *mut u32;
-    // SAFETY: the first 4 bytes is reserved for storing CPU ID.
-    unsafe {
-        bsp_cpu_id_ptr.write(0);
-    }
-
     arch::cpu::local::set_base(bsp_base_va as u64);
 
     has_init::set_true();

ostd/src/cpu/mod.rs

@@ -11,43 +11,85 @@ cfg_if::cfg_if! {
 }
 
 use alloc::vec::Vec;
-use core::sync::atomic::{AtomicU32, Ordering};
 
 use bitvec::{
     prelude::{BitVec, Lsb0},
     slice::IterOnes,
 };
+use local::cpu_local_cell;
+use spin::Once;
 
-use crate::arch::{self, boot::smp::get_num_processors};
+use crate::{
+    arch::boot::smp::get_num_processors, task::DisabledPreemptGuard, trap::DisabledLocalIrqGuard,
+};
 
-/// The number of CPUs. Zero means uninitialized.
-static NUM_CPUS: AtomicU32 = AtomicU32::new(0);
+/// The number of CPUs.
+static NUM_CPUS: Once<u32> = Once::new();
 
 /// Initializes the number of CPUs.
 ///
 /// # Safety
 ///
-/// The caller must ensure that this function is called only once at the
-/// correct time when the number of CPUs is available from the platform.
-pub unsafe fn init() {
+/// The caller must ensure that this function is called only once on the BSP
+/// at the correct time when the number of CPUs is available from the platform.
+pub(crate) unsafe fn init_num_cpus() {
     let num_processors = get_num_processors().unwrap_or(1);
-    NUM_CPUS.store(num_processors, Ordering::Release)
+    NUM_CPUS.call_once(|| num_processors);
+}
+
+/// Initializes the number of the current CPU.
+///
+/// # Safety
+///
+/// The caller must ensure that this function is called only once on the
+/// correct CPU with the correct CPU ID.
+pub(crate) unsafe fn set_this_cpu_id(id: u32) {
+    CURRENT_CPU.store(id);
 }
 
 /// Returns the number of CPUs.
 pub fn num_cpus() -> u32 {
-    let num = NUM_CPUS.load(Ordering::Acquire);
-    debug_assert_ne!(num, 0, "The number of CPUs is not initialized");
-    num
+    debug_assert!(
+        NUM_CPUS.get().is_some(),
+        "The number of CPUs is not initialized"
+    );
+    // SAFETY: The number of CPUs is initialized. The unsafe version is used
+    // to avoid the overhead of the check.
+    unsafe { *NUM_CPUS.get_unchecked() }
 }
 
-/// Returns the ID of this CPU.
+/// A marker trait for guard types that can "pin" the current task to the
+/// current CPU.
 ///
-/// The CPU ID is strategically placed at the beginning of the CPU local storage area.
-pub fn this_cpu() -> u32 {
-    // SAFETY: the cpu ID is stored at the beginning of the cpu local area, provided
-    // by the linker script.
-    unsafe { (arch::cpu::local::get_base() as usize as *mut u32).read() }
+/// Such guard types include [`DisabledLocalIrqGuard`] and
+/// [`DisabledPreemptGuard`]. When such guards exist, the CPU executing the
+/// current task is pinned. So getting the current CPU ID or CPU-local
+/// variables are safe.
+///
+/// # Safety
+///
+/// The implementor must ensure that the current task is pinned to the current
+/// CPU while any one of the instances of the implemented structure exists.
+pub unsafe trait PinCurrentCpu {
+    /// Returns the number of the current CPU.
+    fn current_cpu(&self) -> u32 {
+        let id = CURRENT_CPU.load();
+        debug_assert_ne!(id, u32::MAX, "This CPU is not initialized");
+        id
+    }
+}
+
+// SAFETY: When IRQs are disabled, the task cannot be passively preempted and
+// migrates to another CPU. If the task actively calls `yield`, it will not be
+// successful either.
+unsafe impl PinCurrentCpu for DisabledLocalIrqGuard {}
+// SAFETY: When preemption is disabled, the task cannot be preempted and migrates
+// to another CPU.
+unsafe impl PinCurrentCpu for DisabledPreemptGuard {}
+
+cpu_local_cell! {
+    /// The number of the current CPU.
+    static CURRENT_CPU: u32 = u32::MAX;
 }
 
 /// A subset of all CPUs in the system.