diff --git a/os/arceos/modules/axtask/src/run_queue.rs b/os/arceos/modules/axtask/src/run_queue.rs index a1daea1383..b352f58d0c 100644 --- a/os/arceos/modules/axtask/src/run_queue.rs +++ b/os/arceos/modules/axtask/src/run_queue.rs @@ -895,31 +895,40 @@ impl AxRunQueue { let waking_current_task = current_state == TaskState::Blocked && self.cpu_id == this_cpu_id() && crate::current().ptr_eq(&task); - // If the task is blocked, wait for the task to finish its scheduling process. - // See `unblock_task()` for details. - if current_state == TaskState::Blocked { - // Wait for next task's scheduling process to complete. - // If the owning (remote) CPU is still in the middle of schedule() with - // this task (next task) as prev, wait until it's done referencing the task. - // - // Pairs with the `clear_prev_task_on_cpu()`. - // - // Note: - // 1. This should be placed after the judgement of `TaskState::Blocked,`, - // because the task may have been woken up by other cores. - // 2. This can be placed in the front of `switch_to()` - #[cfg(feature = "smp")] - { - // A scheduler tracepoint or other IRQ-safe notification can wake the - // task that is currently being switched out on this CPU. Waiting for - // `on_cpu` there would wait for the very switch we are still inside. - if !waking_current_task { - while task.on_cpu() { - // Wait for the task to finish its scheduling process. - core::hint::spin_loop(); - } - } + // A blocked task woken here may still be finishing its context + // switch-out on its owning CPU: `on_cpu == true` means its registers + // are not yet fully saved. It must NOT be made runnable (enqueued) + // until `on_cpu` is false, or another CPU could resume it with stale + // registers. Pairs with `clear_prev_task_on_cpu()`. + // + // We must NOT busy-spin on `on_cpu` for a task owned by a *remote* + // CPU (as the old code did): two CPUs each spinning with IRQs off, + // each waiting for the other to reach `clear_prev_task_on_cpu()`, is + // a mutual deadlock (whole-board freeze). Instead, hand the enqueue + // to the owning CPU via a lock-free stash it drains once its context + // is saved. `waking_current_task` (self-wake on this CPU, mid-switch) + // keeps the old inline behavior: this CPU finishes the switch in + // program order when it returns. + #[cfg(feature = "smp")] + if current_state == TaskState::Blocked && !waking_current_task && task.on_cpu() { + // Record where the task must land, then stash a reference for the + // owning CPU to enqueue from `clear_prev_task_on_cpu()`. + task.set_cpu_id(self.cpu_id as _); + task.stash_wake(task.clone()); + // Re-check under the SeqCst handshake. If still on its owning CPU, + // that CPU drains the stash after its switch completes — done. + if task.on_cpu() { + return false; + } + // `on_cpu` cleared meanwhile: the owning CPU may already have + // passed its drain point. Whichever side wins `take_wake` does + // the enqueue exactly once. + if task.take_wake().is_none() { + // Owner won the swap; it enqueues + kicks the target. + return false; } + // We won: the reclaimed reference is dropped here; fall through + // and enqueue our own `task` (its context is now saved). } // TODO: priority #[cfg(feature = "smp")] @@ -1088,7 +1097,8 @@ pub(crate) fn migrate_entry(migrated_task: AxTaskRef) { force_kick_remote_cpu(cpu_id); } -/// Clear the `on_cpu` field of previous task running on this CPU. +/// Clear the `on_cpu` field of the previous task running on this CPU, then +/// complete any cross-core wake that was deferred while it was still `on_cpu`. #[cfg(feature = "smp")] pub(crate) unsafe fn clear_prev_task_on_cpu() { let prev = unsafe { PREV_TASK.current_ref_mut_raw() } @@ -1096,7 +1106,42 @@ pub(crate) unsafe fn clear_prev_task_on_cpu() { .expect("PREV_TASK should have been set by switch_to"); // Safety: prev_task's Arc is still alive on the caller's stack at this point // (switch_to has not yet returned), so the pointer is valid. - unsafe { prev.as_ref() }.set_on_cpu(false); + let prev = unsafe { prev.as_ref() }; + // Publish that the context is fully saved. The SeqCst store pairs with the + // waker's `on_cpu()`/`take_wake()` handshake in `put_task_with_state`. + prev.set_on_cpu(false); + // Drain a wake that raced our switch-out. `take_wake` is the single arbiter: + // if the waker did not reclaim it (it saw `on_cpu` still true), we get the + // owned reference and enqueue it now that the context is saved. + if let Some(task) = prev.take_wake() { + let target = task.cpu_id() as usize; + // Leaf lock: `resched()` already dropped this CPU's scheduler lock before + // `switch_to`, so this takes only the target run queue's lock. + get_run_queue(target) + .scheduler + .lock() + .put_prev_task(task, false); + if target != this_cpu_id() { + // Remote target: ask that CPU to reschedule so it picks the task up + // (and wakes if it is idle in `wait_for_irqs`). + #[cfg(feature = "ipi")] + kick_remote_cpu(target); + } else { + // Local target: `kick_remote_cpu(self)` is a no-op, so the reschedule + // the remote waker's IPI used to deliver here would be lost — the + // task could sit un-run until the next tick, or indefinitely if this + // CPU just switched to `idle` and is about to `wait_for_irqs()`. + // `target == this_cpu_id()` arises when `select_wake_run_queue()` + // falls back to the task's `last_cpu`, which is this owning CPU. + // Request a reschedule on THIS CPU instead: the current task + // (`next_task`, possibly `idle`) is forced to reschedule when the + // switch chain unwinds and its preempt guard is released + // (`current_check_preempt_pending` consumes the flag), mirroring the + // reschedule the IPI path (`request_current_reschedule`) performed. + #[cfg(feature = "preempt")] + crate::current().set_force_resched_pending(true); + } + } } pub(crate) fn init() { let cpu_id = this_cpu_id(); diff --git a/os/arceos/modules/axtask/src/task.rs b/os/arceos/modules/axtask/src/task.rs index ff15ec8c0d..dfa705fb17 100644 --- a/os/arceos/modules/axtask/src/task.rs +++ b/os/arceos/modules/axtask/src/task.rs @@ -1,6 +1,8 @@ use alloc::{boxed::Box, string::String, sync::Arc}; #[cfg(not(feature = "stack-guard-page"))] use core::alloc::Layout; +#[cfg(feature = "smp")] +use core::sync::atomic::AtomicPtr; #[cfg(any( feature = "preempt", all(feature = "stack-guard-page", feature = "smp", feature = "ipi") @@ -96,6 +98,18 @@ pub struct TaskInner { /// Used to indicate whether the task is running on a CPU. #[cfg(feature = "smp")] on_cpu: AtomicBool, + /// One-shot cross-core wake handoff. + /// + /// When a remote CPU wins the `Blocked -> Ready` transition for this task + /// while it is still `on_cpu` (its context not yet fully saved on its owning + /// CPU), the waker must NOT enqueue it — and must not spin on `on_cpu` + /// either (that is the cross-core mutual-wake deadlock). Instead it records + /// the target run-queue in `cpu_id` and stashes an owned reference here; the + /// owning CPU drains it in `clear_prev_task_on_cpu()` once `on_cpu` is false, + /// then enqueues + kicks the target. Holds a `*const AxTask` produced by + /// `Arc::into_raw` (null = empty). See `run_queue::put_task_with_state`. + #[cfg(feature = "smp")] + wake_handoff: AtomicPtr, /// A ticket ID used to identify the timer event. /// Set by `set_timer_ticket()` when creating a timer event in `set_alarm_wakeup()`, @@ -361,6 +375,8 @@ impl TaskInner { cpu_id: AtomicU32::new(0), #[cfg(feature = "smp")] on_cpu: AtomicBool::new(false), + #[cfg(feature = "smp")] + wake_handoff: AtomicPtr::new(core::ptr::null_mut()), #[cfg(feature = "preempt")] need_resched: AtomicBool::new(false), #[cfg(feature = "preempt")] @@ -611,17 +627,52 @@ impl TaskInner { /// while it has not finished its scheduling process. /// The `on_cpu field is set to `true` when the task is preparing to run on a CPU, /// and it is set to `false` when the task has finished its scheduling process in `clear_prev_task_on_cpu()`. + /// + /// `SeqCst` because it participates in a store-before-load (Dekker) handshake + /// with [`Self::stash_wake`]/[`Self::take_wake`] across two distinct atomics + /// (`on_cpu` and `wake_handoff`); Acquire/Release would permit the + /// "both sides observe the other's stale value" lost-wakeup execution. #[cfg(feature = "smp")] #[inline] pub(crate) fn on_cpu(&self) -> bool { - self.on_cpu.load(Ordering::Acquire) + self.on_cpu.load(Ordering::SeqCst) } - /// Sets whether the task is running on a CPU. + /// Sets whether the task is running on a CPU. `SeqCst`, see [`Self::on_cpu`]. #[cfg(feature = "smp")] #[inline] pub(crate) fn set_on_cpu(&self, on_cpu: bool) { - self.on_cpu.store(on_cpu, Ordering::Release) + self.on_cpu.store(on_cpu, Ordering::SeqCst) + } + + /// Stash an owned reference for a deferred cross-core wake (see the + /// `wake_handoff` field). Transfers ownership of `task` into the slot via + /// `Arc::into_raw`. Must be paired with exactly one [`Self::take_wake`]. + #[cfg(feature = "smp")] + #[inline] + pub(crate) fn stash_wake(&self, task: AxTaskRef) { + let ptr = Arc::into_raw(task) as *mut AxTask; + // SeqCst: ordered with the `on_cpu` handshake (see `on_cpu`). + self.wake_handoff.store(ptr, Ordering::SeqCst); + } + + /// Atomically consume a stashed deferred-wake reference, if any. Returns the + /// owned `AxTaskRef` to exactly one caller (the swap is the single arbiter); + /// all other callers get `None`. + #[cfg(feature = "smp")] + #[inline] + pub(crate) fn take_wake(&self) -> Option { + let ptr = self + .wake_handoff + .swap(core::ptr::null_mut(), Ordering::SeqCst); + if ptr.is_null() { + None + } else { + // Safety: `ptr` came from `Arc::into_raw` in `stash_wake`, and the + // swap guarantees a single consumer, so this reconstructs the unique + // owning `Arc` exactly once. + Some(unsafe { Arc::from_raw(ptr as *const AxTask) }) + } } }