Refactor the this_cpu API with PinCurrentCpu

ostd/src/task/mod.rs

@@ -9,6 +9,6 @@ pub mod scheduler;
 mod task;
 
 pub use self::{
-    preempt::{disable_preempt, DisablePreemptGuard},
+    preempt::{disable_preempt, DisabledPreemptGuard},
     task::{AtomicCpuId, Priority, Task, TaskAdapter, TaskContextApi, TaskOptions},
 };

ostd/src/task/preempt/cpu_local.rs

@@ -4,7 +4,7 @@
 //! on a CPU with a single 32-bit, CPU-local integer value.
 //!
 //! * Bits from 0 to 30 represents an unsigned counter called `guard_count`,
-//!   which is the number of `DisablePreemptGuard` instances held by the
+//!   which is the number of `DisabledPreemptGuard` instances held by the
 //!   current CPU;
 //! * Bit 31 is set to `!need_preempt`, where `need_preempt` is a boolean value
 //!   that will be set by the scheduler when it decides that the current task

ostd/src/task/preempt/guard.rs

@@ -3,14 +3,14 @@
 /// A guard for disable preempt.
 #[clippy::has_significant_drop]
 #[must_use]
-pub struct DisablePreemptGuard {
+pub struct DisabledPreemptGuard {
     // This private field prevents user from constructing values of this type directly.
     _private: (),
 }
 
-impl !Send for DisablePreemptGuard {}
+impl !Send for DisabledPreemptGuard {}
 
-impl DisablePreemptGuard {
+impl DisabledPreemptGuard {
     fn new() -> Self {
         super::cpu_local::inc_guard_count();
         Self { _private: () }
@@ -23,13 +23,13 @@ impl DisablePreemptGuard {
     }
 }
 
-impl Drop for DisablePreemptGuard {
+impl Drop for DisabledPreemptGuard {
     fn drop(&mut self) {
         super::cpu_local::dec_guard_count();
     }
 }
 
 /// Disables preemption.
-pub fn disable_preempt() -> DisablePreemptGuard {
-    DisablePreemptGuard::new()
+pub fn disable_preempt() -> DisabledPreemptGuard {
+    DisabledPreemptGuard::new()
 }

ostd/src/task/preempt/mod.rs

@@ -3,4 +3,4 @@
 pub(super) mod cpu_local;
 mod guard;
 
-pub use self::guard::{disable_preempt, DisablePreemptGuard};
+pub use self::guard::{disable_preempt, DisabledPreemptGuard};

ostd/src/task/scheduler/fifo_scheduler.rs

@@ -4,9 +4,9 @@ use alloc::{boxed::Box, collections::VecDeque, sync::Arc, vec::Vec};
 
 use super::{inject_scheduler, EnqueueFlags, LocalRunQueue, Scheduler, UpdateFlags};
 use crate::{
-    cpu::{num_cpus, this_cpu},
+    cpu::{num_cpus, PinCurrentCpu},
     sync::SpinLock,
-    task::{AtomicCpuId, Task},
+    task::{disable_preempt, AtomicCpuId, Task},
 };
 
 pub fn init() {
@@ -61,12 +61,18 @@ impl<T: FifoSchedInfo + Send + Sync> Scheduler<T> for FifoScheduler<T> {
     }
 
     fn local_rq_with(&self, f: &mut dyn FnMut(&dyn LocalRunQueue<T>)) {
-        let local_rq: &FifoRunQueue<T> = &self.rq[this_cpu() as usize].disable_irq().lock();
+        let preempt_guard = disable_preempt();
+        let local_rq: &FifoRunQueue<T> = &self.rq[preempt_guard.current_cpu() as usize]
+            .disable_irq()
+            .lock();
         f(local_rq);
     }
 
     fn local_mut_rq_with(&self, f: &mut dyn FnMut(&mut dyn LocalRunQueue<T>)) {
-        let local_rq: &mut FifoRunQueue<T> = &mut self.rq[this_cpu() as usize].disable_irq().lock();
+        let preempt_guard = disable_preempt();
+        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/mod.rs

@@ -12,7 +12,7 @@ use core::sync::atomic::{AtomicBool, Ordering};
 use spin::Once;
 
 use super::{preempt::cpu_local, processor, task::Task};
-use crate::{arch::timer, cpu::this_cpu, prelude::*};
+use crate::{arch::timer, cpu::PinCurrentCpu, prelude::*, task::disable_preempt};
 
 /// Injects a scheduler implementation into framework.
 ///
@@ -140,8 +140,9 @@ pub(crate) fn unpark_target(runnable: Arc<Task>) {
         .enqueue(runnable, EnqueueFlags::Wake);
     if need_preempt_info.is_some() {
         let cpu_id = need_preempt_info.unwrap();
+        let preempt_guard = disable_preempt();
         // FIXME: send IPI to set remote CPU's need_preempt if needed.
-        if cpu_id == this_cpu() {
+        if cpu_id == preempt_guard.current_cpu() {
             cpu_local::set_need_preempt();
         }
     }
@@ -163,8 +164,9 @@ pub(super) fn run_new_task(runnable: Arc<Task>) {
         .enqueue(runnable, EnqueueFlags::Spawn);
     if need_preempt_info.is_some() {
         let cpu_id = need_preempt_info.unwrap();
+        let preempt_guard = disable_preempt();
         // FIXME: send IPI to set remote CPU's need_preempt if needed.
-        if cpu_id == this_cpu() {
+        if cpu_id == preempt_guard.current_cpu() {
             cpu_local::set_need_preempt();
         }
     }