strat9_kernel/process/
scheduler.rs

1//! Scheduler implementation
2//!
3//! Implements a per-CPU, multi-class, SMP scheduler for Strat9-OS with support
4//! for cooperative and preemptive multitasking.
5//!
6//! ## Scheduler locking contract — mandatory total order
7//!
8//! ```text
9//!   GLOBAL_SCHED_STATE                        (rank 1)
10//!     -> SCHED_IDENTITY write                 (rank 2)
11//!       -> BLOCKED_TASKS                      (rank 3)
12//!         -> LOCAL_SCHEDULERS[cpu]            (rank 4)
13//!           -> Task/Process internal lock     (rank 5)
14//! ```
15//!
16//! `SCHED_IDENTITY` read is observational only:
17//! - do not acquire any scheduler lock while holding it;
18//! - copy the result, release it, then begin a mutating operation.
19//!
20//! ### Forbidden
21//! - Acquiring GLOBAL, IDENTITY, BLOCKED, or another LOCAL while holding LOCAL.
22//! - Holding two LOCAL locks simultaneously.
23//! - Allocating, logging, VFS access, IPI sends, or context switching while
24//!   any scheduler spinlock is held.
25//! - Any transition from Blocked directly to Running (must go through Runnable).
26//!
27//! ### Cross-CPU operations
28//! Use a two-phase protocol: detach under the source local lock, release it,
29//! then attach under the destination local lock.
30//!
31//! ### Task ownership invariant
32//! A non-idle task is in exactly one scheduling ownership state:
33//! New, Runnable(cpu), Running(cpu), Blocked, Zombie, or Reaped.
34//!
35//! ## Canonical state machine
36//!
37//! ```text
38//!   New -> Runnable(cpu)
39//!   Runnable(cpu) -> Running(cpu)     [selection / preempt]
40//!   Running(cpu) -> Runnable(cpu)     [tick / yield / preempt]
41//!   Running(cpu) -> Blocked           [block_current]
42//!   Blocked -> Runnable(cpu)          [wake_task]
43//!   Running/Runnable/Blocked -> Zombie [exit / kill]
44//!   Zombie -> Reaped                  [waitpid / reap]
45//! ```
46//!
47//! Each transition has a single owner (the function performing it) and a
48//! defined transaction: which locks are taken, which containers lose the task,
49//! which containers receive it, and when `taskcpu` becomes visible.
50//!
51//! ## Metrics
52//! All counters (FORCE_RESCHED_HINT, RESCHED_IPI_PENDING, ticks, etc.) are
53//! lock-free atomics and must never be used to drive state transitions.
54
55use super::task::{Pid, Task, TaskId, TaskPriority, TaskState, Tid};
56use crate::{
57    arch::{apic, percpu, restore_flags, save_flags_and_cli, timer, timer::NS_PER_TICK},
58    sync::SpinLock,
59};
60use alloc::{collections::BTreeMap, sync::Arc, vec::Vec};
61use core::sync::atomic::{AtomicBool, AtomicU64, Ordering};
62use spin::RwLock as SpinRwLock;
63
64// ---------------------------------------------------------------------------
65// Scheduler state machine
66// ---------------------------------------------------------------------------
67
68/// Formal scheduling state of a task, encoding ownership unambiguously.
69///
70/// Each state identifies the *exclusive owner* of the task:
71///
72/// | State             | Owner                | Present in                  |
73/// |-------------------|----------------------|-----------------------------|
74/// | New               | creator / global     | alltasks only, no queue     |
75/// | Runnable { cpu }  | LOCAL_SCHEDULERS[cpu]| exactly one class queue     |
76/// | Running { cpu }   | LOCAL_SCHEDULERS[cpu]| exactly currenttask on cpu  |
77/// | Blocked           | BLOCKED_TASKS        | exactly BLOCKED_TASKS       |
78/// | Zombie            | GLOBAL_SCHED_STATE   | zombies + alltasks          |
79/// | Reaped            | nobody               | removed from all structures |
80///
81/// `alltasks` may retain a reference for all states except Reaped, but must
82/// **never** be used to determine schedulability.  The source of truth is the
83/// `SchedState` and its owning container.
84#[derive(Debug, Clone, Copy, PartialEq, Eq)]
85pub enum SchedState {
86    /// Task constructed but not yet visible to the scheduler.
87    New,
88    /// Enqueued on `cpu`'s local run queue, awaiting selection.
89    Runnable { cpu: usize },
90    /// Currently executing on `cpu`.
91    Running { cpu: usize },
92    /// Waiting for an event; entry must exist in `BLOCKED_TASKS`.
93    Blocked,
94    /// Exited; waiting to be reaped by parent's `waitpid`.
95    Zombie,
96    /// Fully removed from all scheduler and identity structures.
97    Reaped,
98}
99
100impl SchedState {
101    /// Returns `true` if the task is in a scheduling-eligible state.
102    #[inline]
103    pub fn is_runnable_like(self) -> bool {
104        matches!(
105            self,
106            SchedState::Runnable { .. } | SchedState::Running { .. }
107        )
108    }
109}
110
111// ---------------------------------------------------------------------------
112// Lockdep debug instrumentation
113// ---------------------------------------------------------------------------
114
115/// Lock ranks matching the total order documented in the module header.
116#[derive(Debug, Clone, Copy, PartialEq, Eq)]
117#[repr(u8)]
118pub(crate) enum LockRank {
119    /// GLOBAL_SCHED_STATE — rank 1.
120    Global = 1,
121    /// SCHED_IDENTITY write — rank 2.
122    IdentityW = 2,
123    /// SCHED_IDENTITY read — rank 2R (must not chain to scheduler locks).
124    IdentityR = 3,
125    /// BLOCKED_TASKS — rank 3.
126    Blocked = 4,
127    /// LOCAL_SCHEDULERS[cpu] — rank 4.
128    Local = 5,
129    /// Task/Process internal — rank 5.
130    TaskInternal = 6,
131}
132
133/// Per-CPU lockdep state, active only under `cfg(debug_assertions)`.
134#[cfg(debug_assertions)]
135#[allow(dead_code)]
136pub(crate) struct LockdepState {
137    depth: usize,
138    held: [HeldLock; 6],
139}
140
141#[cfg(debug_assertions)]
142#[derive(Debug, Clone, Copy)]
143pub(crate) struct HeldLock {
144    rank: LockRank,
145    /// Optional CPU index for LOCAL locks.
146    cpu: Option<usize>,
147    /// Caller return address (for diagnostics).
148    caller: usize,
149}
150
151#[cfg(debug_assertions)]
152impl LockdepState {
153    pub(crate) const fn new() -> Self {
154        Self {
155            depth: 0,
156            held: [HeldLock {
157                rank: LockRank::Global,
158                cpu: None,
159                caller: 0,
160            }; 6],
161        }
162    }
163
164    /// Record a lock acquisition.  Panics in debug if ordering is violated.
165    pub(crate) fn acquire(&mut self, rank: LockRank, cpu: Option<usize>, caller: usize) {
166        if self.depth > 0 {
167            let top = self.held[self.depth - 1];
168            // SCHED_IDENTITY read (rank 3) must not chain to any scheduler lock.
169            if top.rank == LockRank::IdentityR {
170                panic!(
171                    "lockdep: SCHED_IDENTITY read held at depth {}, cannot acquire {:?} \
172                     (read path must not chain to scheduler locks)",
173                    self.depth - 1,
174                    rank
175                );
176            }
177            // LOCAL must never be held while acquiring another LOCAL.
178            if rank == LockRank::Local && cpu.is_some() {
179                for i in 0..self.depth {
180                    if self.held[i].rank == LockRank::Local {
181                        panic!(
182                            "lockdep: two LOCAL locks simultaneously (held {:?} at depth {}, \
183                             acquiring LOCAL[{}] at depth {})",
184                            self.held[i],
185                            i,
186                            cpu.unwrap(),
187                            self.depth
188                        );
189                    }
190                }
191            }
192            // General rank check: new rank must be strictly greater than current top.
193            if (rank as u8) <= (top.rank as u8) {
194                panic!(
195                    "lockdep: lock order violation at depth {}: held {:?}, acquiring {:?}",
196                    self.depth - 1,
197                    top,
198                    rank
199                );
200            }
201        }
202        if self.depth < self.held.len() {
203            self.held[self.depth] = HeldLock { rank, cpu, caller };
204        }
205        self.depth += 1;
206    }
207
208    /// Record a lock release.  Asserts LIFO order.
209    pub(crate) fn release(&mut self, rank: LockRank) {
210        if self.depth == 0 {
211            panic!("lockdep: release underflow for {:?}", rank);
212        }
213        self.depth -= 1;
214        let was = self.held[self.depth];
215        if was.rank != rank {
216            panic!(
217                "lockdep: releasing {:?} but top of stack is {:?} (depth {})",
218                rank, was, self.depth
219            );
220        }
221    }
222
223    /// Assert that no scheduler lock is held.
224    pub(crate) fn assert_no_scheduler_locks(&self) {
225        if self.depth > 0 {
226            panic!(
227                "lockdep: asserting no scheduler locks but depth={}, top={:?}",
228                self.depth,
229                self.held[self.depth - 1]
230            );
231        }
232    }
233
234    /// Assert that a specific rank is currently held.
235    pub(crate) fn assert_held(&self, rank: LockRank) {
236        for i in 0..self.depth {
237            if self.held[i].rank == rank {
238                return;
239            }
240        }
241        panic!(
242            "lockdep: expected {:?} to be held, but depth={}",
243            rank, self.depth
244        );
245    }
246
247    /// Return the current depth.
248    pub(crate) fn current_depth(&self) -> usize {
249        self.depth
250    }
251}
252
253/// Per-CPU lockdep state.  Each CPU tracks its own stack of held scheduler locks.
254/// SAFETY: accessed only from the owning CPU with IRQs disabled (no concurrent access).
255#[cfg(debug_assertions)]
256static mut LOCKDEP: [LockdepState; crate::arch::percpu::MAX_CPUS] =
257    [const { LockdepState::new() }; crate::arch::percpu::MAX_CPUS];
258
259/// Record a lock acquisition in the per-CPU lockdep state.
260#[cfg(debug_assertions)]
261#[inline]
262#[track_caller]
263pub(crate) fn lockdep_acquire(rank: LockRank, cpu: Option<usize>) {
264    let caller = core::panic::Location::caller();
265    let addr = caller as *const core::panic::Location<'static> as usize;
266    unsafe { LOCKDEP[current_cpu_index()].acquire(rank, cpu, addr) };
267}
268
269/// Record a lock release in the per-CPU lockdep state.
270#[cfg(debug_assertions)]
271#[inline]
272pub(crate) fn lockdep_release(rank: LockRank) {
273    unsafe { LOCKDEP[current_cpu_index()].release(rank) };
274}
275
276/// Assert that no scheduler locks are held on this CPU.
277#[cfg(debug_assertions)]
278#[inline]
279pub(crate) fn lockdep_assert_no_locks() {
280    unsafe { LOCKDEP[current_cpu_index()].assert_no_scheduler_locks() };
281}
282
283/// Assert that a specific rank is currently held on this CPU.
284#[cfg(debug_assertions)]
285#[inline]
286pub(crate) fn lockdep_assert_held(rank: LockRank) {
287    unsafe { LOCKDEP[current_cpu_index()].assert_held(rank) };
288}
289
290// No-op stubs for release builds
291#[cfg(not(debug_assertions))]
292#[inline]
293pub(crate) fn lockdep_acquire(_rank: LockRank, _cpu: Option<usize>) {}
294#[cfg(not(debug_assertions))]
295#[inline]
296pub(crate) fn lockdep_release(_rank: LockRank) {}
297#[cfg(not(debug_assertions))]
298#[inline]
299pub(crate) fn lockdep_assert_no_locks() {}
300#[cfg(not(debug_assertions))]
301#[inline]
302pub(crate) fn lockdep_assert_held(_rank: LockRank) {}
303
304/// Per-CPU scheduler tick counters used for CPU usage estimation.
305///
306/// - `CPU_TOTAL_TICKS[cpu]`: all timer ticks observed on `cpu`.
307/// - `CPU_IDLE_TICKS[cpu]`: ticks where the idle task was running on `cpu`.
308///
309/// CPU usage over a time window:
310/// `usage = 1 - (delta_idle / delta_total)`.
311static CPU_TOTAL_TICKS: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
312    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
313static CPU_IDLE_TICKS: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
314    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
315static CPU_RT_RUNTIME_TICKS: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
316    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
317static CPU_FAIR_RUNTIME_TICKS: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
318    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
319static CPU_SWITCH_COUNT: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
320    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
321static CPU_PREEMPT_COUNT: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
322    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
323static CPU_STEAL_IN_COUNT: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
324    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
325static CPU_STEAL_OUT_COUNT: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
326    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
327static CPU_TRY_LOCK_FAIL_COUNT: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
328    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
329static RESCHED_IPI_PENDING: [AtomicBool; crate::arch::percpu::MAX_CPUS] =
330    [const { AtomicBool::new(false) }; crate::arch::percpu::MAX_CPUS];
331static IPI_SEND_TRACE_BUDGET: AtomicU64 = AtomicU64::new(64);
332/// Lock-free per-CPU hint: request a local preemption as soon as maybe_preempt
333/// can observe scheduler state. Written from IRQ paths without touching
334/// `GLOBAL_SCHED_STATE`, consumed under scheduler lock in `maybe_preempt`.
335static FORCE_RESCHED_HINT: [AtomicBool; crate::arch::percpu::MAX_CPUS] =
336    [const { AtomicBool::new(false) }; crate::arch::percpu::MAX_CPUS];
337static LAST_STEAL_TICK: [AtomicU64; crate::arch::percpu::MAX_CPUS] =
338    [const { AtomicU64::new(0) }; crate::arch::percpu::MAX_CPUS];
339/// One-shot flag per CPU: set to true after the first preemption is logged.
340/// Prevents flooding the serial port with a preempt trace on every tick.
341pub(crate) static FIRST_PREEMPT_LOGGED: [AtomicBool; crate::arch::percpu::MAX_CPUS] =
342    [const { AtomicBool::new(false) }; crate::arch::percpu::MAX_CPUS];
343
344const STEAL_IMBALANCE_MIN: usize = 2;
345const STEAL_COOLDOWN_TICKS: u64 = 2;
346
347/// Performs the active cpu count operation.
348#[inline]
349pub(crate) fn active_cpu_count() -> usize {
350    crate::arch::smp::cpu_count()
351        .max(1)
352        .min(crate::arch::percpu::MAX_CPUS)
353}
354
355/// Performs the cpu is valid operation.
356#[inline]
357fn cpu_is_valid(cpu: usize) -> bool {
358    cpu < crate::arch::percpu::MAX_CPUS
359}
360
361#[derive(Clone, Copy)]
362pub struct CpuUsageSnapshot {
363    pub cpu_count: usize,
364    pub total_ticks: [u64; crate::arch::percpu::MAX_CPUS],
365    pub idle_ticks: [u64; crate::arch::percpu::MAX_CPUS],
366}
367
368#[derive(Clone, Copy)]
369pub struct SchedulerMetricsSnapshot {
370    pub cpu_count: usize,
371    pub rt_runtime_ticks: [u64; crate::arch::percpu::MAX_CPUS],
372    pub fair_runtime_ticks: [u64; crate::arch::percpu::MAX_CPUS],
373    pub idle_runtime_ticks: [u64; crate::arch::percpu::MAX_CPUS],
374    pub switch_count: [u64; crate::arch::percpu::MAX_CPUS],
375    pub preempt_count: [u64; crate::arch::percpu::MAX_CPUS],
376    pub steal_in_count: [u64; crate::arch::percpu::MAX_CPUS],
377    pub steal_out_count: [u64; crate::arch::percpu::MAX_CPUS],
378    pub try_lock_fail_count: [u64; crate::arch::percpu::MAX_CPUS],
379    pub deferred_work_raised: [u32; crate::arch::percpu::MAX_CPUS],
380    pub deferred_work_processed: [u32; crate::arch::percpu::MAX_CPUS],
381}
382
383#[derive(Clone, Copy)]
384pub struct SchedulerStateSnapshot {
385    pub initialized: bool,
386    pub boot_phase: u8,
387    pub cpu_count: usize,
388    pub pick_order: [crate::process::sched::SchedClassId; 3],
389    pub steal_order: [crate::process::sched::SchedClassId; 2],
390    pub blocked_tasks: usize,
391    pub current_task: [u64; crate::arch::percpu::MAX_CPUS],
392    pub rq_rt: [usize; crate::arch::percpu::MAX_CPUS],
393    pub rq_fair: [usize; crate::arch::percpu::MAX_CPUS],
394    pub rq_idle: [usize; crate::arch::percpu::MAX_CPUS],
395    pub need_resched: [bool; crate::arch::percpu::MAX_CPUS],
396    pub deferred_work_raised: [u32; crate::arch::percpu::MAX_CPUS],
397    pub deferred_work_processed: [u32; crate::arch::percpu::MAX_CPUS],
398}
399
400/// Performs the cpu usage snapshot operation.
401pub fn cpu_usage_snapshot() -> CpuUsageSnapshot {
402    let cpu_count = active_cpu_count();
403    let mut total_ticks = [0u64; crate::arch::percpu::MAX_CPUS];
404    let mut idle_ticks = [0u64; crate::arch::percpu::MAX_CPUS];
405    for i in 0..cpu_count {
406        total_ticks[i] = CPU_TOTAL_TICKS[i].load(Ordering::Relaxed);
407        idle_ticks[i] = CPU_IDLE_TICKS[i].load(Ordering::Relaxed);
408    }
409    CpuUsageSnapshot {
410        cpu_count,
411        total_ticks,
412        idle_ticks,
413    }
414}
415
416/// Performs the scheduler metrics snapshot operation.
417pub fn scheduler_metrics_snapshot() -> SchedulerMetricsSnapshot {
418    let cpu_count = active_cpu_count();
419    let dw = deferred_work::metrics_snapshot();
420    let mut rt_runtime_ticks = [0u64; crate::arch::percpu::MAX_CPUS];
421    let mut fair_runtime_ticks = [0u64; crate::arch::percpu::MAX_CPUS];
422    let mut idle_runtime_ticks = [0u64; crate::arch::percpu::MAX_CPUS];
423    let mut switch_count = [0u64; crate::arch::percpu::MAX_CPUS];
424    let mut preempt_count = [0u64; crate::arch::percpu::MAX_CPUS];
425    let mut steal_in_count = [0u64; crate::arch::percpu::MAX_CPUS];
426    let mut steal_out_count = [0u64; crate::arch::percpu::MAX_CPUS];
427    let mut try_lock_fail_count = [0u64; crate::arch::percpu::MAX_CPUS];
428    for i in 0..cpu_count {
429        rt_runtime_ticks[i] = CPU_RT_RUNTIME_TICKS[i].load(Ordering::Relaxed);
430        fair_runtime_ticks[i] = CPU_FAIR_RUNTIME_TICKS[i].load(Ordering::Relaxed);
431        idle_runtime_ticks[i] = CPU_IDLE_TICKS[i].load(Ordering::Relaxed);
432        switch_count[i] = CPU_SWITCH_COUNT[i].load(Ordering::Relaxed);
433        preempt_count[i] = CPU_PREEMPT_COUNT[i].load(Ordering::Relaxed);
434        steal_in_count[i] = CPU_STEAL_IN_COUNT[i].load(Ordering::Relaxed);
435        steal_out_count[i] = CPU_STEAL_OUT_COUNT[i].load(Ordering::Relaxed);
436        try_lock_fail_count[i] = CPU_TRY_LOCK_FAIL_COUNT[i].load(Ordering::Relaxed);
437    }
438    SchedulerMetricsSnapshot {
439        cpu_count,
440        rt_runtime_ticks,
441        fair_runtime_ticks,
442        idle_runtime_ticks,
443        switch_count,
444        preempt_count,
445        steal_in_count,
446        steal_out_count,
447        try_lock_fail_count,
448        deferred_work_raised: dw.raised,
449        deferred_work_processed: dw.processed,
450    }
451}
452
453/// Performs the reset scheduler metrics operation.
454pub fn reset_scheduler_metrics() {
455    let cpu_count = active_cpu_count();
456    for i in 0..cpu_count {
457        CPU_RT_RUNTIME_TICKS[i].store(0, Ordering::Relaxed);
458        CPU_FAIR_RUNTIME_TICKS[i].store(0, Ordering::Relaxed);
459        CPU_IDLE_TICKS[i].store(0, Ordering::Relaxed);
460        CPU_SWITCH_COUNT[i].store(0, Ordering::Relaxed);
461        CPU_PREEMPT_COUNT[i].store(0, Ordering::Relaxed);
462        CPU_STEAL_IN_COUNT[i].store(0, Ordering::Relaxed);
463        CPU_STEAL_OUT_COUNT[i].store(0, Ordering::Relaxed);
464        CPU_TRY_LOCK_FAIL_COUNT[i].store(0, Ordering::Relaxed);
465    }
466    deferred_work::reset_metrics();
467}
468
469/// Performs the note try lock fail on cpu operation.
470#[inline]
471pub(crate) fn note_try_lock_fail_on_cpu(cpu: usize) {
472    if cpu_is_valid(cpu) {
473        CPU_TRY_LOCK_FAIL_COUNT[cpu].fetch_add(1, Ordering::Relaxed);
474    }
475}
476
477/// Performs the note try lock fail operation.
478#[inline]
479pub fn note_try_lock_fail() {
480    note_try_lock_fail_on_cpu(current_cpu_index());
481}
482
483// ========== Cross-CPU IPI helpers ===============================================
484
485/// Send a reschedule IPI to `cpu_index`.
486/// No-op if APIC is not initialized, or if `cpu_index` is the current CPU
487/// (the caller already handles the local-CPU case via `yield_cpu`).
488fn send_resched_ipi_to_cpu(cpu_index: usize) {
489    if !cpu_is_valid(cpu_index) {
490        return;
491    }
492    if !apic::is_initialized() {
493        return;
494    }
495    let my_cpu = current_cpu_index();
496    let should_trace = IPI_SEND_TRACE_BUDGET
497        .try_update(Ordering::AcqRel, Ordering::Relaxed, |budget| {
498            budget.checked_sub(1)
499        })
500        .is_ok();
501    if let Some(target_apic) = percpu::apic_id_by_cpu_index(cpu_index) {
502        if let Some(my_apic) = percpu::apic_id_by_cpu_index(my_cpu) {
503            if target_apic != my_apic {
504                if RESCHED_IPI_PENDING[cpu_index].swap(true, Ordering::AcqRel) {
505                    if should_trace {
506                        crate::e9_println!(
507                            "[ipi-send-skip] from_cpu={} to_cpu={} to_apic={:#x} pending=1",
508                            my_cpu,
509                            cpu_index,
510                            target_apic
511                        );
512                    }
513                    return;
514                }
515                if should_trace {
516                    crate::e9_println!(
517                        "[ipi-send] from_cpu={} from_apic={:#x} to_cpu={} to_apic={:#x}",
518                        my_cpu,
519                        my_apic,
520                        cpu_index,
521                        target_apic
522                    );
523                }
524                apic::send_resched_ipi(target_apic);
525            }
526        }
527    }
528}
529
530/// Request a local force-reschedule hint for `cpu`.
531#[inline]
532pub(crate) fn request_force_resched_hint(cpu: usize) {
533    if cpu_is_valid(cpu) {
534        FORCE_RESCHED_HINT[cpu].store(true, Ordering::Release);
535    }
536}
537
538/// Consume and clear the local force-reschedule hint for `cpu`.
539#[inline]
540pub(crate) fn take_force_resched_hint(cpu: usize) -> bool {
541    if cpu_is_valid(cpu) {
542        FORCE_RESCHED_HINT[cpu].swap(false, Ordering::AcqRel)
543    } else {
544        false
545    }
546}
547
548/// Global scheduler state — rank 1 (root) in the total lock order.
549///
550/// This is the root lock.  It may precede all other scheduler locks but must
551/// **never** be acquired from a path that already holds a LOCAL, BLOCKED, or
552/// IDENTITY lock.
553///
554/// Protects: `all_tasks`, `task_cpu`, `zombies`, `wake_deadlines`, and the
555/// global `class_table`.  Per-CPU run queues and current-task tracking live
556/// in `LOCAL_SCHEDULERS` (rank 4).  Blocked tasks are in `BLOCKED_TASKS`
557/// (rank 3).  Identity maps are in `SCHED_IDENTITY` (rank 2).
558pub(crate) static GLOBAL_SCHED_STATE: SpinLock<Option<GlobalSchedState>> = SpinLock::new(None);
559
560/// Returns the scheduler lock address for deadlock tracing.
561pub fn debug_scheduler_lock_addr() -> usize {
562    &GLOBAL_SCHED_STATE as *const _ as usize
563}
564
565/// Global tick counter (safe to increment from interrupt context)
566static TICK_COUNT: AtomicU64 = AtomicU64::new(0);
567/// Verbose scheduler trace switch.
568static SCHED_VERBOSE: AtomicBool = AtomicBool::new(false);
569
570/// Performs the sched trace operation.
571#[inline]
572fn sched_trace(args: core::fmt::Arguments<'_>) {
573    if SCHED_VERBOSE.load(Ordering::Relaxed) {
574        log::debug!("[sched] {}", args);
575    }
576}
577
578/// Information needed to perform a context switch after releasing the lock.
579///
580/// # Safety invariants
581///
582/// `SwitchTarget` contains raw pointers into `Arc<Task>` objects.  All five
583/// invariants below must hold at the moment `do_switch_context` / `switch_context`
584/// reads them.  They are established by `yield_cpu_local` under the LOCAL lock
585/// and consumed before any other CPU can observe the pointed-to memory.
586///
587/// 1. **`old_rsp_ptr`** points to `(*source.context.get()).saved_rsp`.
588///    The `source` Arc<Task> is kept alive by `cpu.current_task` or
589///    `cpu.task_to_requeue` for the duration of the switch; no migration,
590///    reap, or exit can invalidate it while the lock is held.
591///
592/// 2. **`new_rsp_ptr`** points to `(*target.context.get()).saved_rsp`.
593///    Same lifetime guarantee as above: `target` is in `cpu.current_task`.
594///
595/// 3. **`old_fpu_ptr` / `new_fpu_ptr`** point into the FPU state areas of
596///    the source / target tasks respectively.  These areas sit inside the
597///    kernel stack and are valid as long as the owning `Arc<Task>` is alive.
598///    FPU state is never modified concurrently: the switch context owns it
599///    exclusively between the save and restore.
600///
601/// 4. **`old_xcr0` / `new_xcr0`** are the XCR0 masks of source and target,
602///    read atomically from `task.xcr0_mask`.  They are only consumed by the
603///    switch assembly which toggles XCR0 around xsave/xrstor.
604///
605/// 5. **Lifetime**: `SwitchTarget` must not outlive the LOCAL lock that
606///    protected its construction.  It is consumed by `do_switch_context`
607///    immediately after the lock is released, before any allocation, IPI,
608///    or scheduler operation.
609///
610/// # Debug checks
611///
612/// Under `cfg(debug_assertions)`, `prepare_switch_target` validates that
613/// `old_rsp` / `new_rsp` fall within their respective kernel stacks before
614/// returning this struct.
615pub(super) struct SwitchTarget {
616    pub(super) old_rsp_ptr: *mut u64,
617    pub(super) new_rsp_ptr: *const u64,
618    pub(super) old_fpu_ptr: *mut u8,
619    pub(super) new_fpu_ptr: *const u8,
620    pub(super) old_xcr0: u64,
621    pub(super) new_xcr0: u64,
622}
623
624// SAFETY: SwitchTarget is only constructed under LOCAL lock on the owning CPU
625// and consumed by do_switch_context before any other CPU can observe the
626// pointed-to memory.  The Arc<Task> objects are kept alive by current_task /
627// task_to_requeue for the full duration.  No concurrent modification of the
628// pointed-to context or FPU state occurs between construction and consumption.
629unsafe impl Send for SwitchTarget {}
630
631/// Result of a non-blocking wait on child exit.
632pub enum WaitChildResult {
633    Reaped {
634        child: TaskId,
635        pid: Pid,
636        status: i32,
637    },
638    NoChildren,
639    StillRunning,
640}
641
642/// Performs the current cpu index operation.
643fn current_cpu_index() -> usize {
644    crate::arch::percpu::current_cpu_index()
645}
646
647struct PerCpuClassRqSet {
648    real_time: crate::process::sched::real_time::RealTimeClassRq,
649    fair: crate::process::sched::fair::FairClassRq,
650    idle: crate::process::sched::idle::IdleClassRq,
651}
652
653impl PerCpuClassRqSet {
654    /// Creates a new instance.
655    fn new() -> Self {
656        Self {
657            real_time: crate::process::sched::real_time::RealTimeClassRq::new(),
658            fair: crate::process::sched::fair::FairClassRq::new(),
659            idle: crate::process::sched::idle::IdleClassRq::new(),
660        }
661    }
662
663    /// Performs the enqueue operation.
664    fn enqueue(&mut self, class: crate::process::sched::SchedClassId, task: Arc<Task>) {
665        use crate::process::sched::SchedClassRq;
666        match class {
667            crate::process::sched::SchedClassId::Fair => self.fair.enqueue(task),
668            crate::process::sched::SchedClassId::RealTime => self.real_time.enqueue(task),
669            crate::process::sched::SchedClassId::Idle => self.idle.enqueue(task),
670        }
671    }
672
673    /// Performs the len by class operation.
674    fn len_by_class(&self, class: crate::process::sched::SchedClassId) -> usize {
675        use crate::process::sched::SchedClassRq;
676        match class {
677            crate::process::sched::SchedClassId::Fair => self.fair.len(),
678            crate::process::sched::SchedClassId::RealTime => self.real_time.len(),
679            crate::process::sched::SchedClassId::Idle => self.idle.len(),
680        }
681    }
682
683    /// Performs the runnable len operation.
684    fn runnable_len(&self) -> usize {
685        self.len_by_class(crate::process::sched::SchedClassId::RealTime)
686            + self.len_by_class(crate::process::sched::SchedClassId::Fair)
687    }
688
689    /// Performs the pick next by class operation.
690    fn pick_next_by_class(
691        &mut self,
692        class: crate::process::sched::SchedClassId,
693    ) -> Option<Arc<Task>> {
694        use crate::process::sched::SchedClassRq;
695        match class {
696            crate::process::sched::SchedClassId::Fair => self.fair.pick_next(),
697            crate::process::sched::SchedClassId::RealTime => self.real_time.pick_next(),
698            crate::process::sched::SchedClassId::Idle => self.idle.pick_next(),
699        }
700    }
701
702    /// Performs the pick next operation.
703    fn pick_next(&mut self, table: &crate::process::sched::SchedClassTable) -> Option<Arc<Task>> {
704        for class in table.pick_order().iter().copied() {
705            if let Some(task) = self.pick_next_by_class(class) {
706                return Some(task);
707            }
708        }
709        None
710    }
711
712    /// Updates current.
713    fn update_current(
714        &mut self,
715        rt: &crate::process::sched::CurrentRuntime,
716        task: &Task,
717        is_yield: bool,
718        table: &crate::process::sched::SchedClassTable,
719    ) -> bool {
720        use crate::process::sched::SchedClassRq;
721        let should_preempt = match table.class_for_task(task) {
722            crate::process::sched::SchedClassId::Fair => {
723                self.fair.update_current(rt, task, is_yield)
724            }
725            crate::process::sched::SchedClassId::RealTime => {
726                self.real_time.update_current(rt, task, is_yield)
727            }
728            crate::process::sched::SchedClassId::Idle => {
729                self.idle.update_current(rt, task, is_yield)
730            }
731        };
732        // Always preempt idle task if there are other tasks ready
733        let any_ready = !self.real_time.is_empty() || !self.fair.is_empty();
734        should_preempt
735            || (table.class_for_task(task) == crate::process::sched::SchedClassId::Idle
736                && any_ready)
737    }
738
739    /// Performs the remove operation.
740    fn remove(&mut self, task_id: crate::process::TaskId) -> bool {
741        use crate::process::sched::SchedClassRq;
742        self.real_time.remove(task_id) || self.fair.remove(task_id) || self.idle.remove(task_id)
743    }
744
745    /// Called once per timer tick.  Increments wait-time counters for all
746    /// queued tasks across all classes.  Used for Fair starvation detection.
747    fn tick_update_wait(&mut self) {
748        use crate::process::sched::SchedClassRq;
749        self.fair.tick_update_wait();
750    }
751
752    /// Performs the steal candidate operation.
753    fn steal_candidate(
754        &mut self,
755        table: &crate::process::sched::SchedClassTable,
756    ) -> Option<Arc<Task>> {
757        for class in table.steal_order().iter().copied() {
758            if let Some(task) = self.pick_next_by_class(class) {
759                return Some(task);
760            }
761        }
762        None
763    }
764}
765
766/// Per-CPU scheduler state
767struct SchedulerCpu {
768    /// Multi-class priority queues
769    class_rqs: PerCpuClassRqSet,
770    /// Currently running task
771    current_task: Option<Arc<Task>>,
772    /// Current runtime accounting
773    current_runtime: crate::process::sched::CurrentRuntime,
774    /// Idle task to run when no other tasks are ready
775    idle_task: Arc<Task>,
776    /// Task that was just preempted and needs to be re-queued
777    task_to_requeue: Option<Arc<Task>>,
778    /// Task that is dying or blocked, to drop outside the scheduler lock
779    task_to_drop: Option<Arc<Task>>,
780    /// Flag indicating if the current task's time slice has expired
781    need_resched: bool,
782    /// Local copy of the class table for hot-path use without GLOBAL lock.
783    /// Updated atomically when the global class table changes.
784    class_table: crate::process::sched::SchedClassTable,
785}
786
787/// Per-CPU local scheduler locks — rank 4 in the total lock order.
788///
789/// Acquired after `GLOBAL_SCHED_STATE`, `SCHED_IDENTITY`, and `BLOCKED_TASKS`.
790/// **Never** hold two LOCAL locks simultaneously.
791/// Cross-CPU operations use a two-phase protocol: detach under source LOCAL,
792/// release, then attach under destination LOCAL.
793#[allow(dead_code)]
794pub(crate) static LOCAL_SCHEDULERS: [SpinLock<Option<SchedulerCpu>>;
795    crate::arch::percpu::MAX_CPUS] = [const { SpinLock::new(None) }; crate::arch::percpu::MAX_CPUS];
796
797/// Blocked tasks registry — rank 3 in the total lock order.
798///
799/// `BLOCKED_TASKS` must be acquired **after** `GLOBAL_SCHED_STATE` and
800/// `SCHED_IDENTITY` (write), and **before** `LOCAL_SCHEDULERS[cpu]`.
801/// It is **never** acquired from a path that already holds a LOCAL lock.
802///
803/// A task appears here if and only if its `SchedState` is `Blocked`.
804pub(crate) static BLOCKED_TASKS: SpinLock<BTreeMap<TaskId, Arc<Task>>> =
805    SpinLock::new(BTreeMap::new());
806
807/// Identity maps — rank 2 (write) / 2R (read) in the total lock order.
808///
809/// `SCHED_IDENTITY` write is acquired after `GLOBAL_SCHED_STATE` and before
810/// `BLOCKED_TASKS` or `LOCAL_SCHEDULERS[cpu]`.
811///
812/// `SCHED_IDENTITY` read is **observational only**: copy the result, release,
813/// then begin any mutating operation.  No scheduler lock may be acquired
814/// while holding the read guard.
815///
816/// Upgraded to `RwLock` so that concurrent readers (`getpid`, `getpgid`,
817/// `get_task_by_pid`) do not serialize.
818pub(crate) static SCHED_IDENTITY: SpinRwLock<SchedIdentity> = SpinRwLock::new(SchedIdentity::new());
819
820/// Identity maps for the scheduler: PID/TID routing, process groups,
821/// session membership, and parent/child relationships.
822///
823/// Lives behind the `SCHED_IDENTITY` lock, separate from `GLOBAL_SCHED_STATE`
824/// so that syscall lookups (`getpid`, `getpgid`, `setpgid`, `setsid`, etc.)
825/// never contend with fork/exit or block/wake paths.
826pub struct SchedIdentity {
827    /// Map userspace PID -> internal TaskId (process leader in current model).
828    pub pid_to_task: BTreeMap<Pid, TaskId>,
829    /// Map userspace TID -> internal TaskId (fast thread lookup).
830    pub tid_to_task: BTreeMap<Tid, TaskId>,
831    /// Map PID -> process group id.
832    pub pid_to_pgid: BTreeMap<Pid, Pid>,
833    /// Map PID -> session id.
834    pub pid_to_sid: BTreeMap<Pid, Pid>,
835    /// Group membership index: pgid -> task ids.
836    pub pgid_members: BTreeMap<Pid, alloc::vec::Vec<TaskId>>,
837    /// Session membership index: sid -> task ids.
838    pub sid_members: BTreeMap<Pid, alloc::vec::Vec<TaskId>>,
839    /// Parent relationship: child -> parent
840    pub parent_of: BTreeMap<TaskId, TaskId>,
841    /// Children list: parent -> children
842    pub children_of: BTreeMap<TaskId, alloc::vec::Vec<TaskId>>,
843}
844
845impl SchedIdentity {
846    /// Creates a new empty identity registry.
847    pub const fn new() -> Self {
848        Self {
849            pid_to_task: BTreeMap::new(),
850            tid_to_task: BTreeMap::new(),
851            pid_to_pgid: BTreeMap::new(),
852            pid_to_sid: BTreeMap::new(),
853            pgid_members: BTreeMap::new(),
854            sid_members: BTreeMap::new(),
855            parent_of: BTreeMap::new(),
856            children_of: BTreeMap::new(),
857        }
858    }
859}
860
861/// Global task registry : cold path: fork, exit, all_tasks scan.
862///
863/// Lock order: acquire GLOBAL_SCHED_STATE before LOCAL when both are needed.
864/// Per-CPU runqueues and current-task tracking live in `LOCAL_SCHEDULERS`.
865/// Blocked tasks are tracked in `BLOCKED_TASKS` (separate lock).
866/// Identity maps (PID/TID, pgid, sid, parent/child) are in `SCHED_IDENTITY` (separate lock).
867/// This struct holds only data that is accessed by cold paths (fork, exit,
868/// all_tasks scan, zombie management, wake deadlines) and is protected by the
869/// `GLOBAL_SCHED_STATE` lock.
870pub struct GlobalSchedState {
871    /// All tasks in the system (for lookup by TaskId)
872    pub(crate) all_tasks: BTreeMap<TaskId, Arc<Task>>,
873    /// Map TaskId -> CPU index (for wake/resume routing)
874    task_cpu: BTreeMap<TaskId, usize>,
875    /// Deadline -> task ids map for sleeping tasks (ordered wakeups).
876    #[allow(dead_code)]
877    wake_deadlines: BTreeMap<u64, alloc::vec::Vec<TaskId>>,
878    /// Task -> deadline reverse index.
879    #[allow(dead_code)]
880    wake_deadline_of: BTreeMap<TaskId, u64>,
881    /// Zombie exit statuses: child -> (exit_code, pid)
882    zombies: BTreeMap<TaskId, (i32, Pid)>,
883    /// Scheduler class table (pick order, steal order, class metadata)
884    class_table: crate::process::sched::SchedClassTable,
885}
886
887/// Performs the validate task context operation.
888fn validate_task_context(task: &Arc<Task>) -> Result<(), &'static str> {
889    let saved_rsp = unsafe { (*task.context.get()).saved_rsp };
890    let stack_base = task.kernel_stack.virt_base.as_u64();
891    let stack_top = stack_base.saturating_add(task.kernel_stack.size as u64);
892
893    if saved_rsp < stack_base || saved_rsp.saturating_add(56) > stack_top {
894        return Err("saved_rsp outside kernel stack bounds");
895    }
896
897    // ABI alignment: saved_rsp must be 8-byte aligned (x86-64 ABI requirement).
898    if saved_rsp & 7 != 0 {
899        return Err("saved_rsp not 8-byte aligned");
900    }
901
902    // Return IP is at [saved_rsp + 48] in our switch frame layout.
903    let ret_ip = unsafe { core::ptr::read_unaligned((saved_rsp + 48) as *const u64) };
904    if ret_ip == 0 {
905        return Err("null return IP in switch frame");
906    }
907
908    // Return IP must be a canonical userspace or kernel address (not in the
909    // non-canonical hole 0x0000_8000_0000_0000..0xFFFF_7FFF_FFFF_FFFF).
910    let canonical = ret_ip < 0x0000_8000_0000_0000 || ret_ip >= 0xFFFF_8000_0000_0000;
911    if !canonical {
912        return Err("non-canonical return IP");
913    }
914
915    // ExtendedState is embedded in the Arc-owned Task, not in kernel_stack.
916    // The owning Arc keeps the save area alive across the context switch.
917    // Check the alignment required by XSAVE/XRSTOR (also sufficient for FXSAVE).
918    let fpu_ptr = task.fpu_state.get() as usize;
919    if fpu_ptr % core::mem::align_of::<crate::process::task::ExtendedState>() != 0 {
920        return Err("FPU state not aligned for context restore");
921    }
922
923    Ok(())
924}
925
926// ---------------------------------------------------------------------------
927// Debug invariant validator
928// ---------------------------------------------------------------------------
929
930/// Validate scheduler-wide invariants in debug builds.
931///
932/// Uses `try_lock` everywhere so it can be called from `finish_switch` without
933/// blocking.  Silently skips if any lock is contended.
934///
935/// Panics with a diagnostic message on the first invariant violation.
936#[cfg(debug_assertions)]
937#[allow(dead_code)]
938pub(crate) fn validate_scheduler_invariants() {
939    let n = active_cpu_count();
940    let mut seen_current: alloc::collections::BTreeSet<TaskId> =
941        alloc::collections::BTreeSet::new();
942
943    // 1. Each CPU has at most one current_task; no task is current on two CPUs.
944    for cpu_idx in 0..n {
945        let guard = match LOCAL_SCHEDULERS[cpu_idx].try_lock() {
946            Some(g) => g,
947            None => return, // Lock contended, skip validation.
948        };
949        if let Some(ref cpu) = *guard {
950            if let Some(ref current) = cpu.current_task {
951                let tid = current.id;
952                if !Arc::ptr_eq(current, &cpu.idle_task) {
953                    assert!(
954                        seen_current.insert(tid),
955                        "scheduler invariant: task {} is current on multiple CPUs (cpu={})",
956                        tid.as_u64(),
957                        cpu_idx
958                    );
959                }
960            }
961        }
962    }
963
964    // 2. task_cpu consistency + zombie isolation.
965    {
966        let sched_guard = match GLOBAL_SCHED_STATE.try_lock() {
967            Some(g) => g,
968            None => return,
969        };
970        if let Some(ref sched) = *sched_guard {
971            for (tid, &cpu_idx) in sched.task_cpu.iter() {
972                let state = sched.all_tasks.get(tid).map(|t| t.get_state());
973                match state {
974                    Some(TaskState::Ready) | Some(TaskState::Running) => {
975                        assert!(
976                            cpu_idx < n,
977                            "scheduler invariant: task_cpu[{}] = {} but cpu_count = {}",
978                            tid.as_u64(),
979                            cpu_idx,
980                            n
981                        );
982                    }
983                    _ => {}
984                }
985            }
986            for (tid, _) in sched.zombies.iter() {
987                assert!(
988                    !seen_current.contains(tid),
989                    "scheduler invariant: zombie task {} is current on a CPU",
990                    tid.as_u64()
991                );
992            }
993        }
994    }
995}
996
997#[cfg(not(debug_assertions))]
998#[inline]
999pub(crate) fn validate_scheduler_invariants() {}
1000
1001/// Check that `SCHED_IDENTITY` read is not held (observational-only contract).
1002#[cfg(debug_assertions)]
1003#[allow(dead_code)]
1004pub(crate) fn assert_no_identity_read_held() {
1005    lockdep_assert_no_locks();
1006}
1007
1008mod core_impl;
1009pub mod deferred_work;
1010pub mod perf_counters;
1011mod runtime_ops;
1012mod task_ops;
1013mod timer_ops;
1014
1015pub use runtime_ops::*;
1016pub use task_ops::*;
1017pub use timer_ops::*;
1018
1019// Re-export deferred work module for metrics and raise functions.
1020pub use deferred_work::{
1021    has_pending, metrics_snapshot as deferred_work_metrics, process_deferred_work,
1022    raise_deferred_work, raise_tick_deferred_work, reset_metrics as reset_deferred_work_metrics,
1023    DeferredWork, DeferredWorkMetrics,
1024};
strat9_kernel/process/scheduler.rs

strat9_kernel/process/
scheduler.rs
