Use CpuId with validity ensurance rather than u32

ostd/src/task/scheduler/fifo_scheduler.rs

@@ -6,7 +6,7 @@ use super::{
     info::CommonSchedInfo, inject_scheduler, EnqueueFlags, LocalRunQueue, Scheduler, UpdateFlags,
 };
 use crate::{
-    cpu::{num_cpus, PinCurrentCpu},
+    cpu::{num_cpus, CpuId, PinCurrentCpu},
     sync::SpinLock,
     task::{disable_preempt, Task},
 };
@@ -25,7 +25,7 @@ struct FifoScheduler<T: CommonSchedInfo> {
 
 impl<T: CommonSchedInfo> FifoScheduler<T> {
     /// Creates a new instance of `FifoScheduler`.
-    fn new(nr_cpus: u32) -> Self {
+    fn new(nr_cpus: usize) -> Self {
         let mut rq = Vec::new();
         for _ in 0..nr_cpus {
             rq.push(SpinLock::new(FifoRunQueue::new()));
@@ -33,14 +33,14 @@ impl<T: CommonSchedInfo> FifoScheduler<T> {
         Self { rq }
     }
 
-    fn select_cpu(&self) -> u32 {
+    fn select_cpu(&self) -> CpuId {
         // FIXME: adopt more reasonable policy once we fully enable SMP.
-        0
+        CpuId::bsp()
     }
 }
 
 impl<T: CommonSchedInfo + Send + Sync> Scheduler<T> for FifoScheduler<T> {
-    fn enqueue(&self, runnable: Arc<T>, flags: EnqueueFlags) -> Option<u32> {
+    fn enqueue(&self, runnable: Arc<T>, flags: EnqueueFlags) -> Option<CpuId> {
         let mut still_in_rq = false;
         let target_cpu = {
             let mut cpu_id = self.select_cpu();
@@ -53,7 +53,7 @@ impl<T: CommonSchedInfo + Send + Sync> Scheduler<T> for FifoScheduler<T> {
             cpu_id
         };
 
-        let mut rq = self.rq[target_cpu as usize].disable_irq().lock();
+        let mut rq = self.rq[target_cpu.as_usize()].disable_irq().lock();
         if still_in_rq && let Err(_) = runnable.cpu().set_if_is_none(target_cpu) {
             return None;
         }
@@ -64,7 +64,7 @@ impl<T: CommonSchedInfo + Send + Sync> Scheduler<T> for FifoScheduler<T> {
 
     fn local_rq_with(&self, f: &mut dyn FnMut(&dyn LocalRunQueue<T>)) {
         let preempt_guard = disable_preempt();
-        let local_rq: &FifoRunQueue<T> = &self.rq[preempt_guard.current_cpu() as usize]
+        let local_rq: &FifoRunQueue<T> = &self.rq[preempt_guard.current_cpu().as_usize()]
             .disable_irq()
             .lock();
         f(local_rq);
@@ -72,7 +72,7 @@ impl<T: CommonSchedInfo + Send + Sync> Scheduler<T> for FifoScheduler<T> {
 
     fn local_mut_rq_with(&self, f: &mut dyn FnMut(&mut dyn LocalRunQueue<T>)) {
         let preempt_guard = disable_preempt();
-        let local_rq: &mut FifoRunQueue<T> = &mut self.rq[preempt_guard.current_cpu() as usize]
+        let local_rq: &mut FifoRunQueue<T> = &mut self.rq[preempt_guard.current_cpu().as_usize()]
             .disable_irq()
             .lock();
         f(local_rq);

ostd/src/task/scheduler/info.rs

@@ -4,7 +4,7 @@
 
 use core::sync::atomic::{AtomicU32, Ordering};
 
-use crate::task::Task;
+use crate::{cpu::CpuId, task::Task};
 
 /// Fields of a task that OSTD will never touch.
 ///
@@ -28,17 +28,22 @@ impl AtomicCpuId {
     /// An `AtomicCpuId` with `AtomicCpuId::NONE` as its inner value is empty.
     const NONE: u32 = u32::MAX;
 
-    fn new(cpu_id: u32) -> Self {
-        Self(AtomicU32::new(cpu_id))
-    }
-
     /// Sets the inner value of an `AtomicCpuId` if it's empty.
     ///
     /// The return value is a result indicating whether the new value was written
-    /// and containing the previous value.
-    pub fn set_if_is_none(&self, cpu_id: u32) -> core::result::Result<u32, u32> {
+    /// and containing the previous value. If the previous value is empty, it returns
+    /// `Ok(())`. Otherwise, it returns `Err(previous_value)` which the previous
+    /// value is a valid CPU ID.
+    pub fn set_if_is_none(&self, cpu_id: CpuId) -> core::result::Result<(), CpuId> {
         self.0
-            .compare_exchange(Self::NONE, cpu_id, Ordering::Relaxed, Ordering::Relaxed)
+            .compare_exchange(
+                Self::NONE,
+                cpu_id.as_usize() as u32,
+                Ordering::Relaxed,
+                Ordering::Relaxed,
+            )
+            .map(|_| ())
+            .map_err(|prev| (prev as usize).try_into().unwrap())
     }
 
     /// Sets the inner value of an `AtomicCpuId` to `AtomicCpuId::NONE`, i.e. makes
@@ -48,19 +53,19 @@ impl AtomicCpuId {
     }
 
     /// Gets the inner value of an `AtomicCpuId`.
-    pub fn get(&self) -> Option<u32> {
+    pub fn get(&self) -> Option<CpuId> {
         let val = self.0.load(Ordering::Relaxed);
         if val == Self::NONE {
             None
         } else {
-            Some(val)
+            Some((val as usize).try_into().ok()?)
         }
     }
 }
 
 impl Default for AtomicCpuId {
     fn default() -> Self {
-        Self::new(Self::NONE)
+        Self(AtomicU32::new(Self::NONE))
     }
 }
 

ostd/src/task/scheduler/mod.rs

@@ -13,7 +13,12 @@ use core::sync::atomic::{AtomicBool, Ordering};
 use spin::Once;
 
 use super::{preempt::cpu_local, processor, Task};
-use crate::{cpu::PinCurrentCpu, prelude::*, task::disable_preempt, timer};
+use crate::{
+    cpu::{CpuId, PinCurrentCpu},
+    prelude::*,
+    task::disable_preempt,
+    timer,
+};
 
 /// Injects a scheduler implementation into framework.
 ///
@@ -40,7 +45,7 @@ pub trait Scheduler<T = Task>: Sync + Send {
     ///
     /// If the `current` of a CPU needs to be preempted, this method returns the id of
     /// that CPU.
-    fn enqueue(&self, runnable: Arc<T>, flags: EnqueueFlags) -> Option<u32>;
+    fn enqueue(&self, runnable: Arc<T>, flags: EnqueueFlags) -> Option<CpuId>;
 
     /// Gets an immutable access to the local runqueue of the current CPU core.
     fn local_rq_with(&self, f: &mut dyn FnMut(&dyn LocalRunQueue<T>));