diff --git a/ostd/src/task/scheduler/mod.rs b/ostd/src/task/scheduler/mod.rs index 59ea7584d..f898ebf1f 100644 --- a/ostd/src/task/scheduler/mod.rs +++ b/ostd/src/task/scheduler/mod.rs @@ -87,7 +87,8 @@ pub fn inject_scheduler(scheduler: &'static dyn Scheduler) { timer::register_callback(|| { SCHEDULER.get().unwrap().mut_local_rq_with(&mut |local_rq| { - if local_rq.update_current(UpdateFlags::Tick) { + let should_pick_next = local_rq.update_current(UpdateFlags::Tick); + if should_pick_next { cpu_local::set_need_preempt(); } }) @@ -179,7 +180,7 @@ pub trait Scheduler: Sync + Send { /// When OSTD is unsure about whether the current task should or can be replaced, /// it will invoke [`LocalRunQueue::try_pick_next`], the fallible version of `pick_next`. /// -/// ## Intern Working +/// ## Internal Working /// /// To guide scheduler implementers, /// we provide a simplified view of how OSTD interacts with local runqueues _internally_ @@ -322,6 +323,7 @@ pub trait LocalRunQueue { /// /// The return value of this method indicates whether an invocation of `pick_next` should be followed /// to find another task to replace the current one. + #[must_use] fn update_current(&mut self, flags: UpdateFlags) -> bool; /// Picks the next runnable task. @@ -383,15 +385,22 @@ pub(crate) fn might_preempt() { if !cpu_local::should_preempt() { return; } - yield_now(); + reschedule(|local_rq| { + let next_task_opt = local_rq.try_pick_next(); + if let Some(next_task) = next_task_opt { + ReschedAction::SwitchTo(next_task.clone()) + } else { + ReschedAction::DoNothing + } + }) } -/// Blocks the current task unless `has_woken()` returns `true`. +/// Blocks the current task unless `has_unparked()` returns `true`. /// /// Note that this method may return due to spurious wake events. It's the caller's responsibility /// to detect them (if necessary). #[track_caller] -pub(crate) fn park_current(has_woken: F) +pub(crate) fn park_current(has_unparked: F) where F: Fn() -> bool, { @@ -399,32 +408,36 @@ where let mut is_first_try = true; reschedule(|local_rq: &mut dyn LocalRunQueue| { - if is_first_try { - if has_woken() { + let next_task_opt = if is_first_try { + if has_unparked() { return ReschedAction::DoNothing; } + is_first_try = false; - // Note the race conditions: the current task may be woken after the above `has_woken` + // Note the race conditions: the current task may be woken after the above `has_unparked` // check, but before the below `dequeue_current` action, we need to make sure that the // wakeup event isn't lost. // - // Currently, for the FIFO scheduler, `Scheduler::enqueue` will try to lock `local_rq` + // Currently, for the FIFO and CFS scheduler, `Scheduler::enqueue` will try to lock `local_rq` // when the above race condition occurs, so it will wait until we finish calling the // `dequeue_current` method and nothing bad will happen. This may need to be revisited // after more complex schedulers are introduced. - local_rq.update_current(UpdateFlags::Wait); + let should_pick_next = local_rq.update_current(UpdateFlags::Wait); current = local_rq.dequeue_current(); - } + should_pick_next.then(|| local_rq.pick_next()) + } else { + local_rq.try_pick_next() + }; - if let Some(next_task) = local_rq.try_pick_next() { + if let Some(next_task) = next_task_opt { if Arc::ptr_eq(current.as_ref().unwrap(), next_task) { + // The current task has been woken and picked as the next runnable task. return ReschedAction::DoNothing; } return ReschedAction::SwitchTo(next_task.clone()); } - is_first_try = false; ReschedAction::Retry }); } @@ -479,9 +492,19 @@ fn set_need_preempt(cpu_id: CpuId) { /// This should only be called if the current is to exit. #[track_caller] pub(super) fn exit_current() -> ! { + let mut is_first_try = true; + reschedule(|local_rq: &mut dyn LocalRunQueue| { - let _ = local_rq.dequeue_current(); - if let Some(next_task) = local_rq.try_pick_next() { + let next_task_opt = if is_first_try { + is_first_try = false; + let should_pick_next = local_rq.update_current(UpdateFlags::Exit); + let _current = local_rq.dequeue_current(); + should_pick_next.then(|| local_rq.pick_next()) + } else { + local_rq.try_pick_next() + }; + + if let Some(next_task) = next_task_opt { ReschedAction::SwitchTo(next_task.clone()) } else { ReschedAction::Retry @@ -495,8 +518,9 @@ pub(super) fn exit_current() -> ! { #[track_caller] pub(super) fn yield_now() { reschedule(|local_rq| { - local_rq.update_current(UpdateFlags::Yield); - if let Some(next_task) = local_rq.try_pick_next() { + let should_pick_next = local_rq.update_current(UpdateFlags::Yield); + let next_task_opt = should_pick_next.then(|| local_rq.pick_next()); + if let Some(next_task) = next_task_opt { ReschedAction::SwitchTo(next_task.clone()) } else { ReschedAction::DoNothing