fix(axtask): defer cross-core wake instead of spinning on remote on_cpu

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,15 +1097,33 @@ 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() }
         .take()
         .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();

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<AxTask>,
 
     /// 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<AxTaskRef> {
+        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) })
+        }
     }
 }