Refactor code related to CPU local memory

Co-authored-by: Chuandong Li <lichuand@pku.edu.cn>

ostd/src/cpu/cpu_local.rs

@@ -0,0 +1,227 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! CPU local storage.
+//!
+//! This module provides a mechanism to define CPU-local objects.
+//!
+//! This is acheived by placing the CPU-local objects in a special section
+//! `.cpu_local`. The bootstrap processor (BSP) uses the objects linked in this
+//! section, and these objects are copied to dynamically allocated local
+//! storage of each application processors (AP) during the initialization
+//! process.
+//!
+//! Such a mechanism exploits the fact that constant values of non-[`Copy`]
+//! types can be bitwise copied. For example, a [`Option<T>`] object, though
+//! being not [`Copy`], have a constant constructor [`Option::None`] that
+//! produces a value that can be bitwise copied to create a new instance.
+//! [`alloc::sync::Arc`] however, don't have such a constructor, and thus cannot
+//! be directly used as a CPU-local object. Wrapping it in a type that has a
+//! constant constructor, like [`Option<T>`], can make it CPU-local.
+
+use core::ops::Deref;
+
+use crate::{
+    cpu::{get_cpu_local_base, set_cpu_local_base},
+    trap::{disable_local, DisabledLocalIrqGuard},
+};
+
+/// Defines a CPU-local variable.
+///
+/// # Example
+///
+/// ```rust
+/// use crate::cpu_local;
+/// use core::cell::RefCell;
+///
+/// cpu_local! {
+///     static FOO: RefCell<u32> = RefCell::new(1);
+///
+///     #[allow(unused)]
+///     pub static BAR: RefCell<f32> = RefCell::new(1.0);
+/// }
+///
+/// println!("FOO VAL: {:?}", *FOO.borrow());
+/// ```
+#[macro_export]
+macro_rules! cpu_local {
+    ($( $(#[$attr:meta])* $vis:vis static $name:ident: $t:ty = $init:expr; )*) => {
+        $(
+            #[link_section = ".cpu_local"]
+            $(#[$attr])* $vis static $name: $crate::CpuLocal<$t> = {
+                let val = $init;
+                // SAFETY: The CPU local variable instantiated is statically
+                // stored in the special `.cpu_local` section.
+                unsafe {
+                    $crate::CpuLocal::__new(val)
+                }
+            };
+        )*
+    };
+}
+
+/// CPU-local objects.
+///
+/// A CPU-local object only gives you immutable references to the underlying value.
+/// To mutate the value, one can use atomic values (e.g., [`AtomicU32`]) or internally mutable
+/// objects (e.g., [`RefCell`]).
+///
+/// [`AtomicU32`]: core::sync::atomic::AtomicU32
+/// [`RefCell`]: core::cell::RefCell
+pub struct CpuLocal<T>(T);
+
+// SAFETY: At any given time, only one task can access the inner value T
+// of a cpu-local variable even if `T` is not `Sync`.
+unsafe impl<T> Sync for CpuLocal<T> {}
+
+// Prevent valid instances of CpuLocal from being copied to any memory
+// area outside the .cpu_local section.
+impl<T> !Copy for CpuLocal<T> {}
+impl<T> !Clone for CpuLocal<T> {}
+
+// In general, it does not make any sense to send instances of CpuLocal to
+// other tasks as they should live on other CPUs to make sending useful.
+impl<T> !Send for CpuLocal<T> {}
+
+impl<T> CpuLocal<T> {
+    /// Initialize a CPU-local object.
+    ///
+    /// Please do not call this function directly. Instead, use the
+    /// `cpu_local!` macro.
+    ///
+    /// # Safety
+    ///
+    /// The caller should ensure that the object initialized by this
+    /// function resides in the `.cpu_local` section. Otherwise the
+    /// behavior is undefined.
+    #[doc(hidden)]
+    pub const unsafe fn __new(val: T) -> Self {
+        Self(val)
+    }
+
+    /// Get access to the underlying value with IRQs disabled.
+    ///
+    /// By this method, you can borrow a reference to the underlying value
+    /// even if `T` is not `Sync`. Because that it is per-CPU and IRQs are
+    /// disabled, no other running task can access it.
+    pub fn borrow_irq_disabled(&self) -> CpuLocalDerefGuard<'_, T> {
+        CpuLocalDerefGuard {
+            cpu_local: self,
+            _guard: disable_local(),
+        }
+    }
+
+    /// Get access to the underlying value through a raw pointer.
+    ///
+    /// This function calculates the virtual address of the CPU-local object based on the per-
+    /// cpu base address and the offset in the BSP.
+    fn get(&self) -> *const T {
+        let offset = {
+            let bsp_va = self as *const _ as usize;
+            let bsp_base = __cpu_local_start as usize;
+            // The implementation should ensure that the CPU-local object resides in the `.cpu_local`.
+            debug_assert!(bsp_va + core::mem::size_of::<T>() <= __cpu_local_end as usize);
+
+            bsp_va - bsp_base as usize
+        };
+
+        let local_base = get_cpu_local_base() as usize;
+        let local_va = local_base + offset;
+
+        // A sanity check about the alignment.
+        debug_assert_eq!(local_va % core::mem::align_of::<T>(), 0);
+
+        local_va as *mut T
+    }
+}
+
+// Considering a preemptive kernel, a CPU-local object may be dereferenced
+// when another task tries to access it. So, we need to ensure that `T` is
+// `Sync` before allowing it to be dereferenced.
+impl<T: Sync> Deref for CpuLocal<T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: it should be properly initialized before accesses.
+        // And we do not create a mutable reference over it. It is
+        // `Sync` so it can be referenced from this task.
+        unsafe { &*self.get() }
+    }
+}
+
+/// A guard for accessing the CPU-local object.
+///
+/// It ensures that the CPU-local object is accessed with IRQs
+/// disabled. It is created by [`CpuLocal::borrow_irq_disabled`].
+/// Do not hold this guard for a long time.
+#[must_use]
+pub struct CpuLocalDerefGuard<'a, T> {
+    cpu_local: &'a CpuLocal<T>,
+    _guard: DisabledLocalIrqGuard,
+}
+
+impl<T> Deref for CpuLocalDerefGuard<'_, T> {
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: it should be properly initialized before accesses.
+        // And we do not create a mutable reference over it. The IRQs
+        // are disabled so it can be referenced from this task.
+        unsafe { &*self.cpu_local.get() }
+    }
+}
+
+/// Initializes the CPU local data for the bootstrap processor (BSP).
+///
+/// # Safety
+///
+/// This function can only called on the BSP, for once.
+///
+/// It must be guaranteed that the BSP will not access local data before
+/// this function being called, otherwise copying non-constant values
+/// will result in pretty bad undefined behavior.
+pub unsafe fn init_on_bsp() {
+    let start_base_va = __cpu_local_start as usize as u64;
+    set_cpu_local_base(start_base_va);
+}
+
+// These symbols are provided by the linker script.
+extern "C" {
+    fn __cpu_local_start();
+    fn __cpu_local_end();
+}
+
+#[cfg(ktest)]
+mod test {
+    use core::{
+        cell::RefCell,
+        sync::atomic::{AtomicU8, Ordering},
+    };
+
+    use ostd_macros::ktest;
+
+    use super::*;
+
+    #[ktest]
+    fn test_cpu_local() {
+        cpu_local! {
+            static FOO: RefCell<usize> = RefCell::new(1);
+            static BAR: AtomicU8 = AtomicU8::new(3);
+        }
+        for _ in 0..10 {
+            let foo_guard = FOO.borrow_irq_disabled();
+            assert_eq!(*foo_guard.borrow(), 1);
+            *foo_guard.borrow_mut() = 2;
+            drop(foo_guard);
+            for _ in 0..10 {
+                assert_eq!(BAR.load(Ordering::Relaxed), 3);
+                BAR.store(4, Ordering::Relaxed);
+                assert_eq!(BAR.load(Ordering::Relaxed), 4);
+                BAR.store(3, Ordering::Relaxed);
+            }
+            let foo_guard = FOO.borrow_irq_disabled();
+            assert_eq!(*foo_guard.borrow(), 2);
+            *foo_guard.borrow_mut() = 1;
+            drop(foo_guard);
+        }
+    }
+}

ostd/src/cpu/mod.rs

@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: MPL-2.0
+
+//! CPU-related definitions.
+
+pub mod cpu_local;
+
+cfg_if::cfg_if! {
+    if #[cfg(target_arch = "x86_64")]{
+        pub use trapframe::GeneralRegs;
+        pub use crate::arch::x86::cpu::*;
+    }
+}