Refactor ostd::task::processor based on faster CPU-local cells

ostd/src/task/processor.rs

@@ -1,59 +1,40 @@
 // SPDX-License-Identifier: MPL-2.0
 
 use alloc::sync::Arc;
-use core::cell::RefCell;
 
 use super::{
     scheduler::{fetch_task, GLOBAL_SCHEDULER},
     task::{context_switch, TaskContext},
     Task, TaskStatus,
 };
-use crate::{cpu::local::PREEMPT_LOCK_COUNT, cpu_local};
+use crate::{cpu::local::PREEMPT_LOCK_COUNT, cpu_local_cell};
 
-pub struct Processor {
-    current: Option<Arc<Task>>,
-    /// A temporary variable used in [`switch_to_task`] to avoid dropping `current` while running
-    /// as `current`.
-    prev_task: Option<Arc<Task>>,
-    idle_task_ctx: TaskContext,
+cpu_local_cell! {
+    /// The `Arc<Task>` (casted by [`Arc::into_raw`]) that is the current task.
+    static CURRENT_TASK_PTR: *const Task = core::ptr::null();
+    /// The previous task on the processor before switching to the current task.
+    /// It is used for delayed resource release since it would be the current
+    /// task's job to recycle the previous resources.
+    static PREVIOUS_TASK_PTR: *const Task = core::ptr::null();
+    /// An unsafe cell to store the context of the bootstrap code.
+    static BOOTSTRAP_CONTEXT: TaskContext = TaskContext::new();
 }
 
-impl Processor {
-    pub const fn new() -> Self {
-        Self {
-            current: None,
-            prev_task: None,
-            idle_task_ctx: TaskContext::new(),
-        }
+/// Retrieves a reference to the current task running on the processor.
+///
+/// It returns `None` if the function is called in the bootstrap context.
+pub(super) fn current_task() -> Option<Arc<Task>> {
+    let ptr = CURRENT_TASK_PTR.load();
+    if ptr.is_null() {
+        return None;
     }
-    fn get_idle_task_ctx_ptr(&mut self) -> *mut TaskContext {
-        &mut self.idle_task_ctx as *mut _
-    }
-    pub fn take_current(&mut self) -> Option<Arc<Task>> {
-        self.current.take()
-    }
-    pub fn current(&self) -> Option<Arc<Task>> {
-        self.current.as_ref().map(Arc::clone)
-    }
-    pub fn set_current_task(&mut self, task: Arc<Task>) {
-        self.current = Some(task.clone());
-    }
-}
-
-cpu_local! {
-    static PROCESSOR: RefCell<Processor> = RefCell::new(Processor::new());
-}
-
-/// Retrieves the current task running on the processor.
-pub fn current_task() -> Option<Arc<Task>> {
-    PROCESSOR.borrow_irq_disabled().borrow().current()
-}
-
-pub(crate) fn get_idle_task_ctx_ptr() -> *mut TaskContext {
-    PROCESSOR
-        .borrow_irq_disabled()
-        .borrow_mut()
-        .get_idle_task_ctx_ptr()
+    // SAFETY: The pointer is set by `switch_to_task` and is guaranteed to be
+    // built with `Arc::into_raw`.
+    let restored = unsafe { Arc::from_raw(ptr) };
+    // To let the `CURRENT_TASK_PTR` still own the task, we clone and forget it
+    // to increment the reference count.
+    let _ = core::mem::ManuallyDrop::new(restored.clone());
+    Some(restored)
 }
 
 /// Calls this function to switch to other task by using GLOBAL_SCHEDULER
@@ -85,11 +66,11 @@ pub fn preempt(task: &Arc<Task>) {
 
 /// Calls this function to switch to other task
 ///
-/// if current task is none, then it will use the default task context and it will not return to this function again
+/// If current task is none, then it will use the default task context and it
+/// will not return to this function again.
 ///
-/// if current task status is exit, then it will not add to the scheduler
-///
-/// before context switch, current task will switch to the next task
+/// If the current task's status not [`TaskStatus::Runnable`], it will not be
+/// added to the scheduler.
 fn switch_to_task(next_task: Arc<Task>) {
     let preemt_lock_count = PREEMPT_LOCK_COUNT.load();
     if preemt_lock_count != 0 {
@@ -99,23 +80,34 @@ fn switch_to_task(next_task: Arc<Task>) {
         );
     }
 
-    let current_task_ctx_ptr = match current_task() {
-        None => get_idle_task_ctx_ptr(),
-        Some(current_task) => {
-            let ctx_ptr = current_task.ctx().get();
+    let irq_guard = crate::trap::disable_local();
 
-            let mut task_inner = current_task.inner_exclusive_access();
+    let current_task_ptr = CURRENT_TASK_PTR.load();
 
-            debug_assert_ne!(task_inner.task_status, TaskStatus::Sleeping);
-            if task_inner.task_status == TaskStatus::Runnable {
-                drop(task_inner);
-                GLOBAL_SCHEDULER.lock_irq_disabled().enqueue(current_task);
-            } else if task_inner.task_status == TaskStatus::Sleepy {
-                task_inner.task_status = TaskStatus::Sleeping;
-            }
+    let current_task_ctx_ptr = if current_task_ptr.is_null() {
+        // SAFETY: Interrupts are disabled, so the pointer is safe to be fetched.
+        unsafe { BOOTSTRAP_CONTEXT.as_ptr_mut() }
+    } else {
+        // SAFETY: The pointer is not NULL and set as the current task.
+        let cur_task_arc = unsafe {
+            let restored = Arc::from_raw(current_task_ptr);
+            let _ = core::mem::ManuallyDrop::new(restored.clone());
+            restored
+        };
 
-            ctx_ptr
+        let ctx_ptr = cur_task_arc.ctx().get();
+
+        let mut task_inner = cur_task_arc.inner_exclusive_access();
+
+        debug_assert_ne!(task_inner.task_status, TaskStatus::Sleeping);
+        if task_inner.task_status == TaskStatus::Runnable {
+            drop(task_inner);
+            GLOBAL_SCHEDULER.lock().enqueue(cur_task_arc);
+        } else if task_inner.task_status == TaskStatus::Sleepy {
+            task_inner.task_status = TaskStatus::Sleeping;
         }
+
+        ctx_ptr
     };
 
     let next_task_ctx_ptr = next_task.ctx().get().cast_const();
@@ -125,17 +117,22 @@ fn switch_to_task(next_task: Arc<Task>) {
     }
 
     // Change the current task to the next task.
-    {
-        let processor_guard = PROCESSOR.borrow_irq_disabled();
-        let mut processor = processor_guard.borrow_mut();
-
-        // We cannot directly overwrite `current` at this point. Since we are running as `current`,
-        // we must avoid dropping `current`. Otherwise, the kernel stack may be unmapped, leading
-        // to soundness problems.
-        let old_current = processor.current.replace(next_task);
-        processor.prev_task = old_current;
+    //
+    // We cannot directly drop `current` at this point. Since we are running as
+    // `current`, we must avoid dropping `current`. Otherwise, the kernel stack
+    // may be unmapped, leading to instant failure.
+    let old_prev = PREVIOUS_TASK_PTR.load();
+    PREVIOUS_TASK_PTR.store(current_task_ptr);
+    CURRENT_TASK_PTR.store(Arc::into_raw(next_task));
+    // Drop the old-previously running task.
+    if !old_prev.is_null() {
+        // SAFETY: The pointer is set by `switch_to_task` and is guaranteed to be
+        // built with `Arc::into_raw`.
+        drop(unsafe { Arc::from_raw(old_prev) });
     }
 
+    drop(irq_guard);
+
     // SAFETY:
     // 1. `ctx` is only used in `schedule()`. We have exclusive access to both the current task
     //    context and the next task context.