Make current_cpu_racy a method of CpuId

kernel/src/lib.rs

@@ -33,7 +33,7 @@ use kcmdline::KCmdlineArg;
 use ostd::{
     arch::qemu::{exit_qemu, QemuExitCode},
     boot::boot_info,
-    cpu::{CpuId, CpuSet, PinCurrentCpu},
+    cpu::{CpuId, CpuSet},
 };
 use process::{spawn_init_process, Process};
 use sched::SchedPolicy;
@@ -104,22 +104,21 @@ pub fn init() {
 
 fn ap_init() {
     fn ap_idle_thread() {
-        let preempt_guard = ostd::task::disable_preempt();
-        let cpu_id = preempt_guard.current_cpu();
-        drop(preempt_guard);
-        log::info!("Kernel idle thread for CPU #{} started.", cpu_id.as_usize());
+        log::info!(
+            "Kernel idle thread for CPU #{} started.",
+            // No races because `ap_idle_thread` runs on a certain AP.
+            CpuId::current_racy().as_usize(),
+        );
 
         loop {
             crate::thread::Thread::yield_now();
             ostd::cpu::sleep_for_interrupt();
         }
     }
-    let preempt_guard = ostd::task::disable_preempt();
-    let cpu_id = preempt_guard.current_cpu();
-    drop(preempt_guard);
 
     ThreadOptions::new(ap_idle_thread)
-        .cpu_affinity(cpu_id.into())
+        // No races because `ap_init` runs on a certain AP.
+        .cpu_affinity(CpuId::current_racy().into())
         .sched_policy(SchedPolicy::Idle)
         .spawn();
 }

ostd/src/arch/x86/kernel/apic/mod.rs

@@ -39,7 +39,7 @@ static APIC_TYPE: Once<ApicType> = Once::new();
 /// let ticks = apic.timer_current_count();
 /// apic.set_timer_init_count(0);
 /// ```
-pub fn get_or_init(guard: &dyn PinCurrentCpu) -> &(dyn Apic + 'static) {
+pub fn get_or_init(_guard: &dyn PinCurrentCpu) -> &(dyn Apic + 'static) {
     struct ForceSyncSend<T>(T);
 
     // SAFETY: `ForceSyncSend` is `Sync + Send`, but accessing its contained value is unsafe.
@@ -59,7 +59,9 @@ pub fn get_or_init(guard: &dyn PinCurrentCpu) -> &(dyn Apic + 'static) {
         static APIC_INSTANCE: Once<ForceSyncSend<Box<dyn Apic + 'static>>> = Once::new();
     }
 
-    let apic_instance = APIC_INSTANCE.get_on_cpu(guard.current_cpu());
+    // No races due to `_guard`, but use `current_racy` to avoid calling via the vtable.
+    // TODO: Find a better way to make `dyn PinCurrentCpu` easy to use?
+    let apic_instance = APIC_INSTANCE.get_on_cpu(crate::cpu::CpuId::current_racy());
 
     // The APIC instance has already been initialized.
     if let Some(apic) = apic_instance.get() {

ostd/src/cpu/mod.rs

@@ -27,16 +27,34 @@ impl CpuId {
     pub const fn as_usize(self) -> usize {
         self.0 as usize
     }
+
+    /// Returns the ID of the current CPU.
+    ///
+    /// This function is safe to call, but is vulnerable to races. The returned CPU
+    /// ID may be outdated if the task migrates to another CPU.
+    ///
+    /// To ensure that the CPU ID is up-to-date, do it under any guards that
+    /// implement the [`PinCurrentCpu`] trait.
+    pub fn current_racy() -> Self {
+        #[cfg(debug_assertions)]
+        assert!(IS_CURRENT_CPU_INITED.load());
+
+        Self(CURRENT_CPU.load())
+    }
 }
 
+/// The error type returned when converting an out-of-range integer to [`CpuId`].
+#[derive(Debug, Clone, Copy)]
+pub struct CpuIdFromIntError;
+
 impl TryFrom<usize> for CpuId {
-    type Error = &'static str;
+    type Error = CpuIdFromIntError;
 
     fn try_from(value: usize) -> Result<Self, Self::Error> {
         if value < num_cpus() {
             Ok(CpuId(value as u32))
         } else {
-            Err("The given CPU ID is out of range")
+            Err(CpuIdFromIntError)
         }
     }
 }
@@ -123,24 +141,10 @@ unsafe fn set_this_cpu_id(id: u32) {
 pub unsafe trait PinCurrentCpu {
     /// Returns the ID of the current CPU.
     fn current_cpu(&self) -> CpuId {
-        current_cpu_racy()
+        CpuId::current_racy()
     }
 }
 
-/// Returns the ID of the current CPU.
-///
-/// This function is safe to call, but is vulnerable to races. The returned CPU
-/// ID may be outdated if the task migrates to another CPU.
-///
-/// To ensure that the CPU ID is up-to-date, do it under any guards that
-/// implements the [`PinCurrentCpu`] trait.
-pub fn current_cpu_racy() -> CpuId {
-    #[cfg(debug_assertions)]
-    assert!(IS_CURRENT_CPU_INITED.load());
-
-    CpuId(CURRENT_CPU.load())
-}
-
 // SAFETY: A guard that enforces the atomic mode requires disabling any
 // context switching. So naturally, the current task is pinned on the CPU.
 unsafe impl<T: InAtomicMode> PinCurrentCpu for T {}

ostd/src/mm/tlb.rs

@@ -195,8 +195,8 @@ cpu_local! {
 }
 
 fn do_remote_flush() {
-    // No races because we are in IRQs/have disabled preempts.
-    let current_cpu = crate::cpu::current_cpu_racy();
+    // No races because we are in IRQs or have disabled preemption.
+    let current_cpu = crate::cpu::CpuId::current_racy();
 
     let mut new_op_queue = OpsStack::new();
     {

ostd/src/smp.rs

@@ -78,16 +78,15 @@ cpu_local! {
 }
 
 fn do_inter_processor_call(_trapframe: &TrapFrame) {
-    // TODO: in interrupt context, disabling interrupts is not necessary.
-    let preempt_guard = trap::disable_local();
-    let cur_cpu = preempt_guard.current_cpu();
+    // No races because we are in IRQs.
+    let this_cpu_id = crate::cpu::CpuId::current_racy();
 
-    let mut queue = CALL_QUEUES.get_on_cpu(cur_cpu).lock();
+    let mut queue = CALL_QUEUES.get_on_cpu(this_cpu_id).lock();
     while let Some(f) = queue.pop_front() {
         log::trace!(
             "Performing inter-processor call to {:#?} on CPU {:#?}",
             f,
-            cur_cpu
+            this_cpu_id,
         );
         f();
     }