From c74237d2c10754b1341e03f495caafd261cf8b7e Mon Sep 17 00:00:00 2001 From: JosephJoshua Date: Thu, 2 Jul 2026 16:45:05 +0800 Subject: [PATCH 1/2] fix(axtask): defer cross-core wake instead of spinning on remote on_cpu MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit put_task_with_state, when unblocking a Blocked task, busy-spun `while task.on_cpu() { spin_loop() }` to wait for the task to finish being switched out on its owning CPU before enqueuing it (so it is never resumed with unsaved context). On SMP this spins with IRQs disabled: if two CPUs each wake a task the other is mid-switching, each spins on the other's on_cpu forever — neither can reach clear_prev_task_on_cpu to clear it — a whole-board freeze. A single cross-core wake is bounded, so this only bites under concurrent wakes from two CPUs (e.g. cross-core device-completion wakes, or the load balancer), which is why it never reproduces at max_cpu_num=1. Fix: never spin on a remote on_cpu. When the woken task is still on_cpu, record the target run queue and stash an owned reference via a lock-free, SeqCst one-shot slot (wake_handoff), then re-check on_cpu. If it cleared, we enqueue; if not, the owning CPU drains the stash from clear_prev_task_on_cpu() once its context is saved (on_cpu false) and enqueues + kicks the target. The waker takes no lock in the deferred case (only atomics), so no cross-core spin and no IPI-queue cycle can form. A SeqCst store-before-load (Dekker) handshake on on_cpu + wake_handoff makes exactly one side perform the enqueue: no lost wakeup, no double-enqueue, and the enqueue always happens after on_cpu is false, so the context-save invariant the spin protected still holds. Builds aarch64 (smp+ipi); clippy 18/18; fmt clean; QEMU smp4 concurrency regression (test-futex-race, test-fcntl-deadlock-smp, test-sched-family, affinity-bug-sched-affinity-migrate) all pass. --- os/arceos/modules/axtask/src/run_queue.rs | 79 +++++++++++++++-------- os/arceos/modules/axtask/src/task.rs | 57 +++++++++++++++- 2 files changed, 107 insertions(+), 29 deletions(-) diff --git a/os/arceos/modules/axtask/src/run_queue.rs b/os/arceos/modules/axtask/src/run_queue.rs index a1daea1383..f43e25d337 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,24 @@ 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); + #[cfg(feature = "ipi")] + kick_remote_cpu(target); + } } 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) }) + } } } From bf3b71dd4622be7436431ed6138fade0b06955fa Mon Sep 17 00:00:00 2001 From: JosephJoshua Date: Tue, 7 Jul 2026 23:17:16 +0800 Subject: [PATCH 2/2] fix(axtask): request local reschedule when draining a deferred wake to this CPU MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses review: `clear_prev_task_on_cpu` always called `kick_remote_cpu(target)` after draining a deferred wake, but that is a no-op when `target == this_cpu_id()`. This case arises from `select_wake_run_queue()` falling back to the task's `last_cpu` — which is this owning CPU — so the reschedule the remote waker's IPI used to deliver was lost. The re-enqueued task could then sit un-run until the next tick, or indefinitely if this CPU had just switched to `idle` and was about to `wait_for_irqs()`. For the local-target drain, request a reschedule on this CPU instead of the self no-op kick: set `force_resched_pending` on the current task (`next_task`, possibly `idle`). It is consumed by `current_check_preempt_pending` when the switch chain unwinds and the preempt guard is released, mirroring exactly what the IPI path (`request_current_reschedule` -> `force_resched`) performed on the remote CPU. Remote targets still `kick_remote_cpu(target)` as before. Builds aarch64 (smp+ipi+preempt); clippy 18/18; fmt clean; QEMU smp4 concurrency regression passes. --- os/arceos/modules/axtask/src/run_queue.rs | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/os/arceos/modules/axtask/src/run_queue.rs b/os/arceos/modules/axtask/src/run_queue.rs index f43e25d337..b352f58d0c 100644 --- a/os/arceos/modules/axtask/src/run_queue.rs +++ b/os/arceos/modules/axtask/src/run_queue.rs @@ -1121,8 +1121,26 @@ pub(crate) unsafe fn clear_prev_task_on_cpu() { .scheduler .lock() .put_prev_task(task, false); - #[cfg(feature = "ipi")] - kick_remote_cpu(target); + 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() {