Implement CurrentTask

ostd/src/sync/wait.rs

@@ -180,7 +180,7 @@ impl Waiter {
     pub fn new_pair() -> (Self, Arc<Waker>) {
         let waker = Arc::new(Waker {
             has_woken: AtomicBool::new(false),
-            task: Task::current().unwrap(),
+            task: Task::current().unwrap().cloned(),
         });
         let waiter = Self {
             waker: waker.clone(),

ostd/src/task/mod.rs

@@ -11,7 +11,10 @@ mod utils;
 
 use core::{
     any::Any,
+    borrow::Borrow,
     cell::{Cell, SyncUnsafeCell},
+    ops::Deref,
+    ptr::NonNull,
 };
 
 use kernel_stack::KernelStack;
@@ -49,8 +52,11 @@ impl Task {
     /// Gets the current task.
     ///
     /// It returns `None` if the function is called in the bootstrap context.
-    pub fn current() -> Option<Arc<Task>> {
-        current_task()
+    pub fn current() -> Option<CurrentTask> {
+        let current_task = current_task()?;
+
+        // SAFETY: `current_task` is the current task.
+        Some(unsafe { CurrentTask::new(current_task) })
     }
 
     pub(super) fn ctx(&self) -> &SyncUnsafeCell<TaskContext> {
@@ -157,7 +163,7 @@ impl TaskOptions {
         /// all task will entering this function
         /// this function is mean to executing the task_fn in Task
         extern "C" fn kernel_task_entry() -> ! {
-            let current_task = current_task()
+            let current_task = Task::current()
                 .expect("no current task, it should have current task in kernel task entry");
 
             // SAFETY: The `func` field will only be accessed by the current task in the task
@@ -170,7 +176,10 @@ impl TaskOptions {
 
             // Manually drop all the on-stack variables to prevent memory leakage!
             // This is needed because `scheduler::exit_current()` will never return.
-            drop(current_task);
+            //
+            // However, `current_task` _borrows_ the current task without holding
+            // an extra reference count. So we do nothing here.
+
             scheduler::exit_current();
         }
 
@@ -214,6 +223,57 @@ impl TaskOptions {
     }
 }
 
+/// The current task.
+///
+/// This type is not `Send`, so it cannot outlive the current task.
+#[derive(Debug)]
+pub struct CurrentTask(NonNull<Task>);
+
+// The intern `NonNull<Task>` contained by `CurrentTask` implies that `CurrentTask` is `!Send`.
+// But it is still good to do this explicitly because this property is key for soundness.
+impl !Send for CurrentTask {}
+
+impl CurrentTask {
+    /// # Safety
+    ///
+    /// The caller must ensure that `task` is the current task.
+    unsafe fn new(task: NonNull<Task>) -> Self {
+        Self(task)
+    }
+
+    /// Returns a cloned `Arc<Task>`.
+    pub fn cloned(&self) -> Arc<Task> {
+        let ptr = self.0.as_ptr();
+
+        // SAFETY: The current task is always a valid task and it is always contained in an `Arc`.
+        unsafe { Arc::increment_strong_count(ptr) };
+
+        // SAFETY: We've increased the reference count in the current `Arc<Task>` above.
+        unsafe { Arc::from_raw(ptr) }
+    }
+}
+
+impl Deref for CurrentTask {
+    type Target = Task;
+
+    fn deref(&self) -> &Self::Target {
+        // SAFETY: The current task is always a valid task.
+        unsafe { self.0.as_ref() }
+    }
+}
+
+impl AsRef<Task> for CurrentTask {
+    fn as_ref(&self) -> &Task {
+        self
+    }
+}
+
+impl Borrow<Task> for CurrentTask {
+    fn borrow(&self) -> &Task {
+        self
+    }
+}
+
 /// Trait for manipulating the task context.
 pub trait TaskContextApi {
     /// Sets instruction pointer

ostd/src/task/processor.rs

@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: MPL-2.0
 
 use alloc::sync::Arc;
+use core::ptr::NonNull;
 
 use super::{context_switch, Task, TaskContext};
 use crate::cpu_local_cell;
@@ -16,21 +17,11 @@ cpu_local_cell! {
     static BOOTSTRAP_CONTEXT: TaskContext = TaskContext::new();
 }
 
-/// Retrieves a reference to the current task running on the processor.
+/// Returns a pointer 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;
-    }
-    // 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)
+pub(super) fn current_task() -> Option<NonNull<Task>> {
+    NonNull::new(CURRENT_TASK_PTR.load().cast_mut())
 }
 
 /// Calls this function to switch to other task

ostd/src/user.rs

@@ -4,7 +4,7 @@
 
 //! User space.
 
-use crate::{cpu::UserContext, mm::VmSpace, prelude::*, task::Task, trap::TrapFrame};
+use crate::{cpu::UserContext, mm::VmSpace, prelude::*, trap::TrapFrame};
 
 /// A user space.
 ///
@@ -112,7 +112,6 @@ pub trait UserContextApi {
 /// }
 /// ```
 pub struct UserMode<'a> {
-    current: Arc<Task>,
     user_space: &'a Arc<UserSpace>,
     context: UserContext,
 }
@@ -126,7 +125,6 @@ impl<'a> UserMode<'a> {
     /// Creates a new `UserMode`.
     pub fn new(user_space: &'a Arc<UserSpace>) -> Self {
         Self {
-            current: Task::current().unwrap(),
             user_space,
             context: user_space.init_ctx,
         }
@@ -147,7 +145,6 @@ impl<'a> UserMode<'a> {
     where
         F: FnMut() -> bool,
     {
-        debug_assert!(Arc::ptr_eq(&self.current, &Task::current().unwrap()));
         crate::task::atomic_mode::might_sleep();
         self.context.execute(has_kernel_event)
     }