From 1469059888be281fc44510e758a50f4f2d8d8fc6 Mon Sep 17 00:00:00 2001 From: Ruihan Li Date: Sun, 3 Nov 2024 23:11:58 +0800 Subject: [PATCH] Implement `CurrentTask` --- ostd/src/sync/wait.rs | 2 +- ostd/src/task/mod.rs | 68 +++++++++++++++++++++++++++++++++++--- ostd/src/task/processor.rs | 17 +++------- ostd/src/user.rs | 5 +-- 4 files changed, 70 insertions(+), 22 deletions(-) diff --git a/ostd/src/sync/wait.rs b/ostd/src/sync/wait.rs index 94199aa40..dd4d6ce8f 100644 --- a/ostd/src/sync/wait.rs +++ b/ostd/src/sync/wait.rs @@ -180,7 +180,7 @@ impl Waiter { pub fn new_pair() -> (Self, Arc) { 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(), diff --git a/ostd/src/task/mod.rs b/ostd/src/task/mod.rs index bffba7964..1fb81183e 100644 --- a/ostd/src/task/mod.rs +++ b/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> { - current_task() + pub fn current() -> Option { + let current_task = current_task()?; + + // SAFETY: `current_task` is the current task. + Some(unsafe { CurrentTask::new(current_task) }) } pub(super) fn ctx(&self) -> &SyncUnsafeCell { @@ -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); + +// The intern `NonNull` 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) -> Self { + Self(task) + } + + /// Returns a cloned `Arc`. + pub fn cloned(&self) -> Arc { + 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` 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 for CurrentTask { + fn as_ref(&self) -> &Task { + self + } +} + +impl Borrow for CurrentTask { + fn borrow(&self) -> &Task { + self + } +} + /// Trait for manipulating the task context. pub trait TaskContextApi { /// Sets instruction pointer diff --git a/ostd/src/task/processor.rs b/ostd/src/task/processor.rs index 4fe814bac..955e45b2a 100644 --- a/ostd/src/task/processor.rs +++ b/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> { - 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::new(CURRENT_TASK_PTR.load().cast_mut()) } /// Calls this function to switch to other task diff --git a/ostd/src/user.rs b/ostd/src/user.rs index dc7e59249..c516b166d 100644 --- a/ostd/src/user.rs +++ b/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, user_space: &'a Arc, context: UserContext, } @@ -126,7 +125,6 @@ impl<'a> UserMode<'a> { /// Creates a new `UserMode`. pub fn new(user_space: &'a Arc) -> 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) }