strat9_kernel/process/scheduler/

core_impl.rs

use super::{runtime_ops::idle_task_main, *};

/// Create a `SchedulerCpu` for the given CPU index (creates its idle task).
pub(super) fn create_cpu_scheduler(cpu_idx: usize) -> SchedulerCpu {
    crate::serial_println!(
        "[trace][sched] create_cpu_scheduler cpu={} create idle begin",
        cpu_idx
    );
    let idle_task = Task::new_kernel_task(idle_task_main, "idle", TaskPriority::Idle)
        .expect("Failed to create idle task");
    crate::serial_println!(
        "[trace][sched] create_cpu_scheduler cpu={} create idle done id={}",
        cpu_idx,
        idle_task.id.as_u64()
    );
    idle_task.set_sched_policy(crate::process::sched::SchedPolicy::Idle);
    let mut class_rqs = PerCpuClassRqSet::new();
    class_rqs.enqueue(crate::process::sched::SchedClassId::Idle, idle_task.clone());
    SchedulerCpu {
        class_rqs,
        current_task: None,
        current_runtime: crate::process::sched::CurrentRuntime::new(),
        idle_task,
        task_to_requeue: None,
        task_to_drop: None,
        need_resched: false,
        class_table: crate::process::sched::SchedClassTable::default(),
    }
}

impl GlobalSchedState {
    /// Create a new global scheduler state (no per-CPU runqueues : those live in LOCAL_SCHEDULERS).
    pub fn new() -> Self {
        crate::serial_println!("[trace][sched] GlobalSchedState::new enter");
        GlobalSchedState {
            all_tasks: BTreeMap::new(),
            task_cpu: BTreeMap::new(),
            wake_deadlines: BTreeMap::new(),
            wake_deadline_of: BTreeMap::new(),
            zombies: BTreeMap::new(),
            class_table: crate::process::sched::SchedClassTable::default(),
        }
    }

    /// Performs the member add operation.
    pub(crate) fn member_add(
        map: &mut BTreeMap<Pid, alloc::vec::Vec<TaskId>>,
        key: Pid,
        task_id: TaskId,
    ) {
        let members = map.entry(key).or_default();
        if !members.iter().any(|id| *id == task_id) {
            members.push(task_id);
        }
    }

    /// Performs the member remove operation.
    pub(crate) fn member_remove(
        map: &mut BTreeMap<Pid, alloc::vec::Vec<TaskId>>,
        key: Pid,
        task_id: TaskId,
    ) {
        let mut clear = false;
        if let Some(members) = map.get_mut(&key) {
            members.retain(|id| *id != task_id);
            clear = members.is_empty();
        }
        if clear {
            map.remove(&key);
        }
    }

    /// Performs the register identity locked operation.
    pub(crate) fn register_identity_locked(identity: &mut SchedIdentity, task: &Arc<Task>) {
        let task_id = task.id;
        let pid = task.pid;
        let pgid = task.pgid.load(Ordering::Relaxed);
        let sid = task.sid.load(Ordering::Relaxed);
        crate::serial_println!(
            "[trace][sched] register_identity enter tid={} pid={} pgid={} sid={}",
            task_id.as_u64(),
            pid,
            pgid,
            sid
        );
        identity.pid_to_pgid.insert(pid, pgid);
        crate::serial_println!(
            "[trace][sched] register_identity pid_to_pgid inserted pid={}",
            pid
        );
        identity.pid_to_sid.insert(pid, sid);
        crate::serial_println!(
            "[trace][sched] register_identity pid_to_sid inserted pid={}",
            pid
        );
        Self::member_add(&mut identity.pgid_members, pgid, task_id);
        Self::member_add(&mut identity.sid_members, sid, task_id);
        crate::serial_println!(
            "[trace][sched] register_identity done tid={}",
            task_id.as_u64()
        );
    }

    /// Performs the unregister identity locked operation.
    pub(crate) fn unregister_identity_locked(
        identity: &mut SchedIdentity,
        task_id: TaskId,
        pid: Pid,
        tid: Tid,
    ) {
        identity.pid_to_task.remove(&pid);
        identity.tid_to_task.remove(&tid);
        if let Some(pgid) = identity.pid_to_pgid.remove(&pid) {
            Self::member_remove(&mut identity.pgid_members, pgid, task_id);
        }
        if let Some(sid) = identity.pid_to_sid.remove(&pid) {
            Self::member_remove(&mut identity.sid_members, sid, task_id);
        }
    }

    /// Add a task to the scheduler
    pub fn add_task(&mut self, task: Arc<Task>) -> Option<usize> {
        let cpu_index = self.select_cpu_for_task(&task);
        self.add_task_on_cpu(task, cpu_index)
    }

    /// Performs the add task with parent operation.
    pub fn add_task_with_parent(&mut self, task: Arc<Task>, parent: TaskId) -> Option<usize> {
        let child = task.id;
        let cpu_index = self.select_cpu_for_task(&task);
        let ipi = self.add_task_on_cpu(task, cpu_index);
        {
            let mut identity = SCHED_IDENTITY.write();
            identity.parent_of.insert(child, parent);
            identity.children_of.entry(parent).or_default().push(child);
        }
        ipi
    }

    /// Performs the add task on cpu operation.
    fn add_task_on_cpu(&mut self, task: Arc<Task>, cpu_index: usize) -> Option<usize> {
        let task_id = task.id;
        crate::serial_println!(
            "[trace][sched] add_task_on_cpu enter tid={} cpu={}",
            task_id.as_u64(),
            cpu_index
        );
        task.set_state(TaskState::Ready);

        self.insert_all_task_locked(task_id, task.clone());
        self.task_cpu.insert(task_id, cpu_index);
        task.home_cpu
            .store(cpu_index, core::sync::atomic::Ordering::Relaxed);

        // Lock order: IDENTITY write (rank 2). Caller holds GLOBAL (rank 1).
        lockdep_acquire(LockRank::IdentityW, None);
        {
            let mut identity = SCHED_IDENTITY.write();
            identity.pid_to_task.insert(task.pid, task_id);
            identity.tid_to_task.insert(task.tid, task_id);
            Self::register_identity_locked(&mut identity, &task);
        }
        lockdep_release(LockRank::IdentityW);

        // Lock order: LOCAL (rank 4). Caller holds GLOBAL (rank 1).
        lockdep_acquire(LockRank::Local, Some(cpu_index));
        {
            let class = self.class_table.class_for_task(&task);
            if let Some(ref mut local_cpu) = *LOCAL_SCHEDULERS[cpu_index].lock() {
                local_cpu.class_rqs.enqueue(class, task);
                local_cpu.need_resched = true;
            }
        }
        lockdep_release(LockRank::Local);

        sched_trace(format_args!(
            "enqueue task={} cpu={}",
            task_id.as_u64(),
            cpu_index
        ));
        if cpu_index != current_cpu_index() {
            Some(cpu_index)
        } else {
            None
        }
    }

    pub(super) fn insert_all_task_locked(&mut self, task_id: TaskId, task: Arc<Task>) {
        assert_eq!(
            task.id,
            task_id,
            "scheduler corruption: insert_all_task_locked task.id={} != task_id={}",
            task.id.as_u64(),
            task_id.as_u64()
        );
        if self.all_tasks.contains_key(&task_id) {
            unsafe {
                crate::arch::serial::putc(b'D');
            }
            crate::serial_force_println!(
                "[RACE] insert_all_task_locked: duplicate tid={} all_tasks={}",
                task_id.as_u64(),
                self.all_tasks.len(),
            );
            panic!(
                "scheduler corruption: duplicate insert_all_task_locked tid={}",
                task_id.as_u64()
            );
        }
        self.all_tasks.insert(task_id, task);
    }

    pub(super) fn remove_all_task_locked(&mut self, task_id: TaskId) -> Option<Arc<Task>> {
        self.all_tasks.remove(&task_id)
    }

    /// Performs the clear task wake deadline locked operation.
    pub fn clear_task_wake_deadline_locked(&mut self, id: TaskId) -> bool {
        if let Some(task) = self.all_tasks.get(&id) {
            task.wake_deadline_ns.store(0, Ordering::Relaxed);
            true
        } else {
            false
        }
    }

    /// Sets task wake deadline locked.
    pub fn set_task_wake_deadline_locked(&mut self, id: TaskId, deadline: u64) -> bool {
        if deadline == 0 {
            return self.clear_task_wake_deadline_locked(id);
        }
        if let Some(task) = self.all_tasks.get(&id) {
            task.wake_deadline_ns.store(deadline, Ordering::Relaxed);
            true
        } else {
            false
        }
    }

    /// Performs the wake task locked operation.
    ///
    /// Returns `(was_woken, ipi_cpu)`. The caller must send a resched IPI to
    /// `ipi_cpu` after releasing the scheduler lock.
    ///
    /// NOTE: `blocked_tasks` lives in the separate `BLOCKED_TASKS` lock.
    /// This method handles waking the task directly if it is already in `BLOCKED_TASKS`,
    /// or setting `wake_pending` as a fallback if the task is still transitioning to Blocked.
    pub fn wake_task_locked(&mut self, id: TaskId) -> (bool, Option<usize>) {
        self.clear_task_wake_deadline_locked(id);

        // Lock order: BLOCKED (rank 3) → LOCAL (rank 4). Caller holds GLOBAL (rank 1).
        let mut woken = false;
        let mut ipi_cpu = None;
        {
            lockdep_acquire(LockRank::Blocked, None);
            let mut blocked = super::BLOCKED_TASKS.lock();
            if let Some(task) = blocked.remove(&id) {
                task.set_state(TaskState::Ready);

                // --- CPU placement: prefer last_cpu (cache warmth), then home_cpu ---
                let last = task.last_cpu.load(Ordering::Relaxed);
                let home = task.home_cpu.load(Ordering::Relaxed);
                let n = active_cpu_count();

                // Try last_cpu first: valid, online, and reasonably loaded.
                let cpu_index = if last < n {
                    let last_ok = {
                        lockdep_acquire(LockRank::Local, Some(last));
                        let ok = LOCAL_SCHEDULERS[last]
                            .lock()
                            .as_ref()
                            .map(|c| c.class_rqs.runnable_len() <= 2)
                            .unwrap_or(false);
                        lockdep_release(LockRank::Local);
                        ok
                    };
                    if last_ok {
                        last
                    } else if home < n {
                        home
                    } else {
                        0
                    }
                } else if home < n {
                    home
                } else {
                    0
                };

                let class = {
                    use crate::process::sched::SchedClassId;
                    match task.sched_policy() {
                        crate::process::sched::SchedPolicy::RealTimeRR { .. }
                        | crate::process::sched::SchedPolicy::RealTimeFifo { .. } => {
                            SchedClassId::RealTime
                        }
                        crate::process::sched::SchedPolicy::Fair(_) => SchedClassId::Fair,
                        crate::process::sched::SchedPolicy::Idle => SchedClassId::Idle,
                    }
                };

                lockdep_acquire(LockRank::Local, Some(cpu_index));
                if let Some(ref mut local_cpu) = *super::LOCAL_SCHEDULERS[cpu_index].lock() {
                    local_cpu.class_rqs.enqueue(class, task.clone());
                    local_cpu.need_resched = true;
                }
                lockdep_release(LockRank::Local);

                ipi_cpu = if cpu_index != current_cpu_index() {
                    Some(cpu_index)
                } else {
                    None
                };
                woken = true;
            }
            lockdep_release(LockRank::Blocked);
            drop(blocked);
        }

        if woken {
            return (true, ipi_cpu);
        }

        // Fallback: task not yet in BLOCKED_TASKS (still transitioning to
        // Blocked). Set wake_pending so block_current_task skips blocking.
        if let Some(task) = self.all_tasks.get(&id) {
            task.wake_pending
                .store(true, core::sync::atomic::Ordering::Release);
            (true, None)
        } else {
            (false, None)
        }
    }

    /// Attempts to reap child locked.
    /// If `target` is `Some(tid)`, only reaps that child; otherwise, reaps any child.
    /// Returns `WaitChildResult` indicating the outcome.
    /// Must be called with the scheduler lock held.
    ///
    pub fn try_reap_child_locked(
        &mut self,
        parent: TaskId,
        target: Option<TaskId>,
    ) -> WaitChildResult {
        // First, check children under SCHED_IDENTITY lock.
        let target_is_child = {
            let identity = SCHED_IDENTITY.read();
            let Some(children_view) = identity.children_of.get(&parent) else {
                return WaitChildResult::NoChildren;
            };

            if children_view.is_empty() {
                return WaitChildResult::NoChildren;
            }

            let target_is_child = if let Some(target_id) = target {
                children_view.iter().any(|&id| id == target_id)
            } else {
                true
            };
            target_is_child
        };

        if !target_is_child {
            return WaitChildResult::NoChildren;
        }

        // Find the zombie child : re-check children under SCHED_IDENTITY.
        let zombie = {
            let identity = SCHED_IDENTITY.read();
            let children = match identity.children_of.get(&parent) {
                Some(c) => c.clone(),
                None => return WaitChildResult::NoChildren,
            };
            children
                .iter()
                .copied()
                .find(|id| target.map_or(true, |t| t == *id) && self.zombies.contains_key(id))
        };

        if let Some(child) = zombie {
            let (status, child_pid) = self.zombies.remove(&child).unwrap_or((0, 0));
            // Remove from all_tasks now so that pick_next_task (if it races with
            // reaping) will see was_registered=false and skip cleanup_task_resources.
            let reaped_task = self.remove_all_task_locked(child);
            if let Some(task) = reaped_task.as_ref() {
                super::task_ops::cleanup_task_resources(task);
            }
            let child_tid = reaped_task.as_ref().map(|t| t.tid);
            if child_pid != 0 {
                if let Some(tid) = child_tid {
                    let mut identity = SCHED_IDENTITY.write();
                    Self::unregister_identity_locked(&mut identity, child, child_pid, tid);
                }
            }
            {
                let mut identity = SCHED_IDENTITY.write();
                if let Some(children) = identity.children_of.get_mut(&parent) {
                    children.retain(|&id| id != child);
                    if children.is_empty() {
                        identity.children_of.remove(&parent);
                    }
                }
                identity.parent_of.remove(&child);
            }
            return WaitChildResult::Reaped {
                child,
                pid: child_pid,
                status,
            };
        }

        WaitChildResult::StillRunning
    }

    /// Select the least-loaded CPU for a newly created task.
    ///
    /// Called with `GLOBAL_SCHED_STATE` held.  Acquires each LOCAL sequentially.
    /// Lock order: GLOBAL (held) → LOCAL[cpu] (rank 4).  Never two LOCALs simultaneously.
    ///
    /// Placement strategy:
    /// 1. Early boot (all idle): CPU 0.
    /// 2. Soft affinity: prefer CPUs in the task's `affinity_mask` (from silo config).
    ///    Tie-break with `last_cpu` (cache warmth).  If no affinity-eligible CPU
    ///    has strictly lower load, fall back to unrestricted search.
    /// 3. Unrestricted: pick the globally least-loaded CPU.
    fn select_cpu_for_task(&self, task: &Task) -> usize {
        let n = active_cpu_count();
        let all_idle = (0..n).all(|i| {
            lockdep_acquire(LockRank::Local, Some(i));
            let result = LOCAL_SCHEDULERS[i]
                .lock()
                .as_ref()
                .map(|cpu| cpu.current_task.is_none())
                .unwrap_or(true);
            lockdep_release(LockRank::Local);
            result
        });
        if all_idle {
            crate::serial_println!("[trace][sched] select_cpu_for_task early-boot best=0");
            return 0;
        }

        let affinity = task.affinity_mask.load(Ordering::Relaxed);
        let last = task.last_cpu.load(Ordering::Relaxed);
        let has_affinity = affinity != 0;

        // Pass 1: find least-loaded CPU among affinity-eligible CPUs.
        let mut best_aff = None;
        let mut best_aff_load = usize::MAX;
        let mut best_aff_last = None; // best among those matching last_cpu
        let mut best_aff_last_load = usize::MAX;

        for idx in 0..n {
            let eligible = if has_affinity {
                (affinity & (1u64 << idx)) != 0
            } else {
                true
            };
            if !eligible {
                continue;
            }
            lockdep_acquire(LockRank::Local, Some(idx));
            let load = {
                let guard = LOCAL_SCHEDULERS[idx].lock();
                if let Some(ref cpu) = *guard {
                    let mut l = cpu.class_rqs.runnable_len();
                    if let Some(current) = cpu.current_task.as_ref() {
                        if self.class_table.class_for_task(current)
                            != crate::process::sched::SchedClassId::Idle
                        {
                            l += 1;
                        }
                    }
                    l
                } else {
                    0
                }
            };
            lockdep_release(LockRank::Local);

            if load < best_aff_load {
                best_aff = Some(idx);
                best_aff_load = load;
            }
            // Prefer last_cpu when loads are equal (cache warmth).
            if idx == last && load <= best_aff_last_load {
                best_aff_last = Some(idx);
                best_aff_last_load = load;
            }
        }

        // If we have an affinity-eligible candidate, prefer it.
        // Use last_cpu tie-breaker when loads are equal.
        if let Some(aff_cpu) = best_aff {
            // If last_cpu is also eligible and has equal or lower load, prefer it.
            if let Some(lc) = best_aff_last {
                if best_aff_last_load <= best_aff_load {
                    crate::serial_println!(
                        "[trace][sched] select_cpu_for_task affinity+last_cpu cpu={} load={}",
                        lc,
                        best_aff_last_load
                    );
                    return lc;
                }
            }
            crate::serial_println!(
                "[trace][sched] select_cpu_for_task affinity cpu={} load={}",
                aff_cpu,
                best_aff_load
            );
            return aff_cpu;
        }

        // Pass 2 (fallback): unrestricted least-loaded.
        let mut best = 0usize;
        let mut best_load = usize::MAX;
        for idx in 0..n {
            lockdep_acquire(LockRank::Local, Some(idx));
            let load = {
                let guard = LOCAL_SCHEDULERS[idx].lock();
                if let Some(ref cpu) = *guard {
                    let mut l = cpu.class_rqs.runnable_len();
                    if let Some(current) = cpu.current_task.as_ref() {
                        if self.class_table.class_for_task(current)
                            != crate::process::sched::SchedClassId::Idle
                        {
                            l += 1;
                        }
                    }
                    l
                } else {
                    0
                }
            };
            lockdep_release(LockRank::Local);
            if load < best_load {
                best = idx;
                best_load = load;
            }
        }
        crate::serial_println!(
            "[trace][sched] select_cpu_for_task fallback best={} load={}",
            best,
            best_load
        );
        best
    }

    /// Migrate all ready tasks to match the current class table.
    ///
    /// Called with `GLOBAL_SCHED_STATE` held.  Acquires each `LOCAL_SCHEDULERS[cpu]`
    /// sequentially (never two at once).  Lock order: GLOBAL (held) → LOCAL[cpu] (rank 4).
    pub fn migrate_ready_tasks_for_new_class_table(&mut self) {
        let mut ready: Vec<(TaskId, Arc<Task>, usize)> = Vec::new();
        for (id, task) in self.all_tasks.iter() {
            let state = task.get_state();
            if state != TaskState::Ready {
                continue;
            }
            let cpu = self.task_cpu.get(id).copied().unwrap_or(0);
            ready.push((*id, task.clone(), cpu));
        }

        for (id, task, cpu_idx) in ready {
            lockdep_acquire(LockRank::Local, Some(cpu_idx));
            let mut guard = LOCAL_SCHEDULERS[cpu_idx].lock();
            let Some(ref mut cpu) = *guard else {
                lockdep_release(LockRank::Local);
                continue;
            };
            if cpu.class_rqs.remove(id) {
                let class = self.class_table.class_for_task(&task);
                cpu.class_rqs.enqueue(class, task);
                cpu.need_resched = true;
            }
            lockdep_release(LockRank::Local);
        }
    }
}

//  Per-CPU hot-path helpers
//
// These functions operate primarily on `SchedulerCpu` (acquired via
// `LOCAL_SCHEDULERS[cpu_index]`).  Most never touch the global `SCHEDULER`
// lock.  The one exception is `steal_task_local`, which does a **non-blocking**
// `GLOBAL_SCHED_STATE.try_lock_no_irqsave()` to update `task_cpu` after a successful
// steal.  This is an intentional lock-order inversion (LOCAL held, then GLOBAL
// attempted) that is safe because the try-lock never blocks : if GLOBAL_SCHED_STATE is
// contended, we simply skip stealing.
//
// Lock order for steal: own LOCAL held => try_lock GLOBAL => try_lock sibling
// LOCALs.  Never blocking-wait, so no deadlock possible.

/// Steal a task from the busiest sibling CPU using per-CPU LOCAL locks.
///
/// Called with `cpu` borrowed from `LOCAL_SCHEDULERS[cpu_index]` (our own
/// LOCAL lock already held). Uses `try_lock_no_irqsave` on sibling entries :
/// if a sibling or the global scheduler state is contended, we skip stealing
/// rather than waiting.
pub(super) fn steal_task_local(cpu: &mut SchedulerCpu, cpu_index: usize) -> Option<Arc<Task>> {
    let now_tick = TICK_COUNT.load(Ordering::Relaxed);
    if now_tick < LAST_STEAL_TICK[cpu_index].load(Ordering::Relaxed) + STEAL_COOLDOWN_TICKS {
        return None;
    }
    lockdep_assert_held(LockRank::Local);

    // Best-effort only: if a cold path is holding the global scheduler, skip
    // stealing instead of blocking the hot path.
    let mut scheduler = match GLOBAL_SCHED_STATE.try_lock_no_irqsave() {
        Some(g) => g,
        None => return None,
    };
    lockdep_acquire(LockRank::Global, None);
    let result = steal_task_inner(&mut scheduler, cpu, cpu_index, now_tick);
    lockdep_release(LockRank::Global);
    result
}

/// Inner implementation of steal, separated for clean lockdep release.
fn steal_task_inner(
    scheduler: &mut Option<GlobalSchedState>,
    cpu: &mut SchedulerCpu,
    cpu_index: usize,
    now_tick: u64,
) -> Option<Arc<Task>> {
    let sched = scheduler.as_mut()?;

    let n = active_cpu_count();
    let my_load = cpu.class_rqs.runnable_len();

    let mut best_cpu = None;
    let mut best_load = 0usize;

    for i in 0..n {
        if i == cpu_index {
            continue;
        }
        // try_lock_no_irqsave: returns immediately if contended (no deadlock).
        if let Some(guard) = LOCAL_SCHEDULERS[i].try_lock_no_irqsave() {
            if let Some(ref sib) = *guard {
                let load = sib.class_rqs.runnable_len();
                if load > best_load {
                    best_load = load;
                    best_cpu = Some(i);
                }
            }
        }
    }

    if best_load < my_load.saturating_add(STEAL_IMBALANCE_MIN) {
        return None;
    }
    let steal_from = best_cpu?;

    // Re-acquire the sibling lock to perform the steal.
    if let Some(mut guard) = LOCAL_SCHEDULERS[steal_from].try_lock_no_irqsave() {
        if let Some(ref mut sib) = *guard {
            if sib.class_rqs.runnable_len() < 2 {
                return None;
            }
            if let Some(task) = sib.class_rqs.steal_candidate(&sib.class_table) {
                // Soft affinity check: skip if the stealing CPU is not in the
                // task's affinity mask.  Re-enqueue on the source to avoid
                // dropping the task.
                let affinity = task.affinity_mask.load(Ordering::Relaxed);
                if affinity != 0 && (affinity & (1u64 << cpu_index)) == 0 {
                    // Re-enqueue on source CPU.
                    let class = sib.class_table.class_for_task(&task);
                    sib.class_rqs.enqueue(class, task);
                    return None;
                }
                sched.task_cpu.insert(task.id, cpu_index);
                task.home_cpu
                    .store(cpu_index, core::sync::atomic::Ordering::Relaxed);
                if cpu_is_valid(cpu_index) {
                    CPU_STEAL_IN_COUNT[cpu_index].fetch_add(1, Ordering::Relaxed);
                }
                if cpu_is_valid(steal_from) {
                    CPU_STEAL_OUT_COUNT[steal_from].fetch_add(1, Ordering::Relaxed);
                }
                LAST_STEAL_TICK[cpu_index].store(now_tick, Ordering::Relaxed);
                return Some(task);
            }
        }
    }
    None
}

/// Pick the next task using only per-CPU LOCAL state.
///
/// Handles current task disposition (re-queue, drop-for-cleanup, or ignore if
/// Blocked), then picks from the local class_rqs, falls back to work-stealing,
/// and finally returns the idle task.
///
/// **Dead tasks**: if the current task is Dead, it goes into `task_to_drop`.
/// Global map cleanup (`all_tasks`, `task_cpu`, etc.) must have been performed
/// by the caller (e.g., `exit_current_task`) BEFORE reaching this point.
pub(super) fn pick_next_task_local(cpu: &mut SchedulerCpu, cpu_index: usize) -> Arc<Task> {
    // Step 1: dispose of the current task.
    if let Some(task) = cpu.current_task.take() {
        match task.get_state() {
            TaskState::Running => {
                task.set_state(TaskState::Ready);
                if !Arc::ptr_eq(&task, &cpu.idle_task) {
                    // Defer re-queue to finish_switch (not yet safe to enqueue :
                    // another CPU could steal it before our context is saved).
                    cpu.task_to_requeue = Some(task);
                }
            }
            TaskState::Dead => {
                // Global maps already cleaned by exit_current_task / kill_task.
                // Defer the Arc drop so KernelStack::drop => buddy_alloc runs
                // outside any lock.
                cpu.task_to_drop = Some(task);
            }
            TaskState::Blocked | TaskState::Ready => {
                // Blocked: moved to blocked_tasks by block_current_task : do nothing.
                // Ready: shouldn't normally occur for current_task; safe to ignore.
            }
        }
    }

    // Step 2: pick from local class_rqs.
    unsafe {
        core::arch::asm!("out 0xe9, al", in("al") b'O', options(nomem, nostack));
    }
    let next = if let Some(next) = cpu.class_rqs.pick_next(&cpu.class_table) {
        unsafe {
            core::arch::asm!("out 0xe9, al", in("al") b'J', options(nomem, nostack));
        }
        next
    } else {
        // Step 3: local queues empty — try work-stealing before idle.
        unsafe {
            core::arch::asm!("out 0xe9, al", in("al") b'S', options(nomem, nostack));
        }
        if let Some(stolen) = steal_task_local(cpu, cpu_index) {
            unsafe {
                core::arch::asm!("out 0xe9, al", in("al") b's', options(nomem, nostack));
            }
            stolen
        } else {
            unsafe {
                core::arch::asm!("out 0xe9, al", in("al") b'j', options(nomem, nostack));
            }
            // Step 4: idle fallback.
            cpu.idle_task.clone()
        }
    };

    next.set_state(TaskState::Running);
    next.last_cpu.store(cpu_index, Ordering::Relaxed);
    cpu.current_task = Some(next.clone());
    cpu.current_runtime = crate::process::sched::CurrentRuntime::new();
    next
}

/// Prepare a LOCAL-only context switch.
///
/// Updates the TSS, SYSCALL RSP, and CR3 for the next task. Returns the
/// raw pointer pair needed by `do_switch_context`.
///
/// Returns `None` if there is no task to switch to (same task or invalid context).
pub(super) fn yield_cpu_local(cpu: &mut SchedulerCpu, cpu_index: usize) -> Option<SwitchTarget> {
    let current = cpu.current_task.as_ref()?.clone();

    let next = pick_next_task_local(cpu, cpu_index);

    if Arc::ptr_eq(&current, &next) {
        return None;
    }
    if cpu_is_valid(cpu_index) {
        CPU_SWITCH_COUNT[cpu_index].fetch_add(1, Ordering::Relaxed);
    }

    if let Err(e) = validate_task_context(&next) {
        let bad_rsp = unsafe { (*next.context.get()).saved_rsp };
        let stk_base = next.kernel_stack.virt_base.as_u64();
        let stk_top = stk_base + next.kernel_stack.size as u64;
        crate::serial_println!(
            "[sched-local] WARN: invalid ctx task='{}' id={} cpu={}: {} \
             rsp={:#x} stack=[{:#x}..{:#x}] : restoring current",
            next.name,
            next.id.as_u64(),
            cpu_index,
            e,
            bad_rsp,
            stk_base,
            stk_top,
        );

        // Restore invariants: undo what pick_next_task_local mutated.
        let is_idle = Arc::ptr_eq(&next, &cpu.idle_task);
        drop(cpu.task_to_drop.take());
        if let Some(prev) = cpu.task_to_requeue.take() {
            prev.set_state(TaskState::Running);
            cpu.current_task = Some(prev);
        } else {
            current.set_state(TaskState::Running);
            cpu.current_task = Some(current.clone());
        }
        if !is_idle {
            next.set_state(TaskState::Ready);
            let class = cpu.class_table.class_for_task(&next);
            cpu.class_rqs.enqueue(class, next);
        }
        return None;
    }

    // Update TSS.rsp0 and SYSCALL kernel RSP for the new task.
    let stack_top = next.kernel_stack.virt_base.as_u64() + next.kernel_stack.size as u64;
    crate::arch::tss::set_kernel_stack(crate::arch::xshim::VirtAddr::new(stack_top));
    crate::arch::syscall::set_kernel_rsp(stack_top);

    // Switch CR3 if the new task has a different address space.
    // SAFETY: The new task's address space has a valid PML4 with the kernel half mapped.
    unsafe {
        next.process.address_space_arc().switch_to();
    }

    Some(SwitchTarget {
        old_rsp_ptr: unsafe { &raw mut (*current.context.get()).saved_rsp },
        new_rsp_ptr: unsafe { &raw const (*next.context.get()).saved_rsp },
        old_fpu_ptr: current.fpu_state.get() as *mut u8,
        new_fpu_ptr: next.fpu_state.get() as *const u8,
        old_xcr0: current
            .xcr0_mask
            .load(core::sync::atomic::Ordering::Relaxed),
        new_xcr0: next.xcr0_mask.load(core::sync::atomic::Ordering::Relaxed),
    })
}

/// Post-switch cleanup using only LOCAL state: re-enqueue the previous task
/// and optionally extract the task-to-drop for deferred deallocation.
pub(super) fn drain_post_switch_local(
    cpu: &mut SchedulerCpu,
    take_drop: bool,
) -> Option<Arc<Task>> {
    let task_to_drop = if take_drop {
        cpu.task_to_drop.take()
    } else {
        None
    };
    if let Some(task) = cpu.task_to_requeue.take() {
        let class = cpu.class_table.class_for_task(&task);
        cpu.class_rqs.enqueue(class, task);
    }
    task_to_drop
}