Skip to main content

rivet/preempt/
mod.rs

1//! Preemptive tier: tasks with dedicated stacks, real priority preemption.
2//!
3//! Unlike `#[rivet::task]` (cooperative — only yields at `.await`), a
4//! preemptive task can be suspended by the timer tick at *any* point and
5//! resumed later from exactly there, because its full execution context
6//! (registers + program counter + stack pointer) is saved/restored on
7//! every switch. This is what lets a genuinely higher-priority task
8//! interrupt a lower-priority one that never calls anything cooperative.
9//!
10//! All switching — tick-driven preemption *and* voluntary yields (blocking
11//! on a mutex, explicit `port::arch::request_reschedule()`) — goes through the same
12//! interrupt/trap path (software interrupt on RISC-V, PendSV on Cortex-M).
13//! There's no separate "synchronous context switch function call" — the
14//! arch trap handler saves the interrupted task's full context, asks
15//! [`on_tick`] which task to resume, and returns to that task's saved
16//! context. This matches how real embedded RTOS ports (FreeRTOS, etc.)
17//! implement it, and keeps there being exactly one code path that has to
18//! be correct instead of two.
19//!
20//! The cooperative async executor still exists — it runs as an ordinary
21//! preemptive task at the lowest priority (see [`crate::init`]), so any
22//! real preemptive task immediately preempts it, and it fills otherwise-idle
23//! CPU time with async work.
24
25pub mod lifecycle;
26pub mod mutex;
27pub mod sched;
28pub mod stack_pool;
29pub mod tcb;
30
31pub use mutex::{PriorityMutex, PriorityMutexGuard};
32pub use tcb::TaskState;
33
34use crate::sync::atomic::Ordering;
35
36/// Statically-sized, correctly-aligned stack storage for a preemptive task.
37#[repr(C, align(16))]
38pub struct Stack<const SIZE: usize>(pub [u8; SIZE]);
39
40impl<const SIZE: usize> Stack<SIZE> {
41    pub const fn new() -> Self {
42        Self([0; SIZE])
43    }
44}
45
46impl<const SIZE: usize> Default for Stack<SIZE> {
47    fn default() -> Self {
48        Self::new()
49    }
50}
51
52/// Implementation helper for [`macro@spawn_ptask`]: allocate the stack from
53/// the pool (or use the provided fallback on host builds) and spawn.
54#[doc(hidden)]
55pub fn spawn_ptask_impl<T: 'static + Send, A: 'static, F: Fn() -> &'static mut [u8]>(
56    stack_size: usize,
57    priority: u8,
58    entry: fn(&'static A) -> T,
59    arg: &'static A,
60    #[allow(unused_variables)] // embedded: the pool is authoritative
61    fallback: F,
62) -> Result<TaskHandle, SpawnError> {
63    let stack = match crate::preempt::stack_pool::alloc_stack(stack_size) {
64        Some(s) => s,
65        None => {
66            // On a real board the pool is authoritative: a fallback stack
67            // outside `.task_stacks` would silently bypass the MPU/PMP
68            // guards (plan.md §4.3). The host test backend (no
69            // linker-provided pool at all) uses the per-invocation static
70            // fallback instead.
71            #[cfg(not(feature = "host-port"))]
72            return Err(SpawnError::StackPoolFull);
73            #[cfg(feature = "host-port")]
74            fallback()
75        }
76    };
77    // SAFETY: the pool slice is exclusively owned by the new task for its
78    // lifetime; the fallback slice has the same contract (see the macro).
79    unsafe { spawn(stack, priority, entry, arg) }
80}
81
82/// Handle to a spawned preemptive task: the registry slot id plus the
83/// slot's generation counter, so stale handles can be detected after the
84/// slot was recycled (plan.md §4.3 / §5.1).
85#[derive(Debug, Clone, Copy, PartialEq, Eq)]
86pub struct TaskHandle {
87    pub id: u16,
88    pub generation: u32,
89}
90
91pub use lifecycle::JoinError;
92
93impl TaskHandle {
94    /// Whether the handle still refers to the same task (slot not recycled).
95    pub fn is_valid(&self) -> bool {
96        tcb::get(self.id as usize)
97            .map(|t| t.generation.load(Ordering::Acquire) == self.generation)
98            .unwrap_or(false)
99    }
100
101    /// Ask the task to stop cooperatively (plan.md §5.4): sets the
102    /// stop-requested flag, which the task polls via
103    /// [`lifecycle::should_stop`]. Returns false if the handle is stale.
104    pub fn request_stop(&self) -> bool {
105        match tcb::get(self.id as usize) {
106            Some(t) if t.generation.load(Ordering::Acquire) == self.generation => {
107                t.stop_requested.store(true, Ordering::Release);
108                true
109            }
110            _ => false,
111        }
112    }
113
114    /// Block until the task's entry returns, then recover its result
115    /// (plan.md §5.2/§5.3). See [`JoinError`].
116    pub fn join<T: 'static + Send>(&self) -> Result<T, JoinError> {
117        lifecycle::join_task::<T>(self)
118    }
119
120    /// Configure this task's period (plan.md Phase 11): the task itself
121    /// calls [`crate::deadlines::wait_period`] once per iteration to block
122    /// until the next boundary. `0` disables periodic waiting. No-op on a
123    /// stale handle.
124    pub fn set_period_us(&self, period_us: u32) {
125        if self.is_valid() {
126            crate::deadlines::set_period_us(self.id as usize, period_us);
127        }
128    }
129
130    /// Configure this task's per-period CPU budget in microseconds
131    /// (plan.md Phase 11, estimated via
132    /// [`crate::exec_time::estimate_us_from_cycles`]). Exceeding it inside
133    /// one period raises [`crate::fault::FaultKind::BudgetExceeded`]
134    /// through the normal fault policy. `0` disables enforcement. No-op on
135    /// a stale handle. Meaningless without also calling
136    /// [`Self::set_period_us`] — the budget window resets at each period
137    /// boundary.
138    pub fn set_budget_us(&self, budget_us: u32) {
139        if self.is_valid() {
140            crate::deadlines::set_budget_us(self.id as usize, budget_us);
141        }
142    }
143
144    /// Release the task's slot and stack for reuse (plan.md §5.4). The
145    /// task must have exited (`join` returned) or be a task other than the
146    /// current one that is blocked/suspended — despawning a *running*
147    /// task is rejected. Returns false for a stale handle.
148    pub fn despawn(&self) -> bool {
149        let Some(t) = tcb::get(self.id as usize) else {
150            return false;
151        };
152        if t.generation.load(Ordering::Acquire) != self.generation {
153            return false;
154        }
155        if sched::current() == Some(self.id as usize) {
156            return false; // cannot despawn the running task
157        }
158        if !t.used.load(Ordering::Acquire) {
159            return false;
160        }
161        // See `PriorityMutexGuard::drop`'s doc for why `ready_remove`
162        // needs a critical section: it isn't a single atomic RMW against
163        // `READY_BITMAP`+`QUEUES` together, so a tick interleaving with
164        // it can observe/leave torn scheduler state.
165        crate::critical::enter(|| crate::preempt::sched::ready_remove(self.id as usize));
166
167        // Drop any stored result, then reset the slot (state READY so the
168        // next `register` claim CAS can succeed; `used=false` publishes).
169        let drop_fn = t.result_drop.load(Ordering::Acquire);
170        if drop_fn != 0 {
171            // SAFETY: `drop_fn` was registered by `spawn` as
172            // `drop_in_place_erased::<T> as *const () as usize` for the
173            // exact T stored in the result buffer.
174            let f: fn(*mut u8) = unsafe { core::mem::transmute(drop_fn) };
175            f(t.result_buf.get() as *mut u8);
176            t.result_drop.store(0, Ordering::Release);
177        }
178        t.state.store(tcb::READY, Ordering::Release);
179        t.exited.store(false, Ordering::Release);
180        t.result_size.store(0, Ordering::Release);
181        t.stop_requested.store(false, Ordering::Release);
182        t.used.store(false, Ordering::Release);
183
184        // Release the stack back to the pool (refilled with 0xAA).
185        let stack = t.stack_info();
186        if let Some((base, size)) = stack {
187            if base != 0 && size != 0 {
188                // SAFETY: the pool slice was given to this task at spawn
189                // and is now unused; `release_stack` refills and recycles it.
190                let slice: &'static mut [u8] =
191                    unsafe { core::slice::from_raw_parts_mut(base as *mut u8, size) };
192                crate::preempt::stack_pool::release_stack(slice);
193            }
194        }
195        t.stack_base.store(0, Ordering::Release);
196        t.stack_size.store(0, Ordering::Release);
197        true
198    }
199
200    /// Suspend the task (plan.md §5.5): a READY task moves to SUSPENDED and
201    /// is skipped by the scheduler until [`TaskHandle::resume`]. Returns
202    /// false for a stale handle or a task that isn't currently READY.
203    pub fn pause(&self) -> bool {
204        let Some(t) = tcb::get(self.id as usize) else {
205            return false;
206        };
207        if t.generation.load(Ordering::Acquire) != self.generation {
208            return false;
209        }
210        // The CAS and `ready_remove` commit as one step (see
211        // `PriorityMutexGuard::drop`'s doc for why) — otherwise a tick
212        // between them could dispatch this task (still queued) even
213        // though its state already reads `SUSPENDED`.
214        crate::critical::enter(|| {
215            let was_ready = t
216                .state
217                .compare_exchange(
218                    tcb::READY,
219                    tcb::SUSPENDED,
220                    Ordering::AcqRel,
221                    Ordering::Acquire,
222                )
223                .is_ok();
224            if was_ready {
225                crate::preempt::sched::ready_remove(self.id as usize);
226            }
227            was_ready
228        })
229    }
230
231    /// Resume a suspended task (plan.md §5.5). Returns false for a stale
232    /// handle or a task that isn't SUSPENDED.
233    pub fn resume(&self) -> bool {
234        let Some(t) = tcb::get(self.id as usize) else {
235            return false;
236        };
237        if t.generation.load(Ordering::Acquire) != self.generation {
238            return false;
239        }
240        // See `pause`'s identical reasoning.
241        crate::critical::enter(|| {
242            let was_suspended = t
243                .state
244                .compare_exchange(
245                    tcb::SUSPENDED,
246                    tcb::READY,
247                    Ordering::AcqRel,
248                    Ordering::Acquire,
249                )
250                .is_ok();
251            if was_suspended {
252                crate::preempt::sched::ready_add(self.id as usize);
253            }
254            was_suspended
255        })
256    }
257}
258
259/// Errors from [`spawn`] / [`macro@spawn_ptask`] — fixed-size resources
260/// degrade with a typed error, never a silent drop or panic (plan.md §4.3).
261#[derive(Debug, Clone, Copy, PartialEq, Eq)]
262pub enum SpawnError {
263    /// The task registry is full ([`tcb::MAX_PTASKS`] slots in use).
264    RegistryFull,
265    /// The task-stack pool is exhausted.
266    StackPoolFull,
267}
268
269/// Spawn a preemptive task with its own stack.
270///
271/// `entry` receives `arg` (a `'static` reference — no heap allocation
272/// needed since the argument just needs to outlive the task, and tasks
273/// never exit). Returns the assigned task id, or `None` if the task
274/// registry ([`tcb::MAX_PTASKS`]) is full.
275///
276/// # Safety
277/// `stack` must not be shared with any other task and must remain valid
278/// (i.e. `'static`) for as long as the task runs.
279pub unsafe fn spawn<T: 'static + Send, A: 'static>(
280    stack: &'static mut [u8],
281    priority: u8,
282    entry: fn(&'static A) -> T,
283    arg: &'static A,
284) -> Result<TaskHandle, SpawnError> {
285    assert!(
286        stack.len() >= crate::port::arch::min_task_stack(),
287        "rivet: task stack too small: {} bytes < arch minimum {} (context-switch frame + entry trampoline; plan.md §2.7)",
288        stack.len(),
289        crate::port::arch::min_task_stack()
290    );
291    // Fill with a known pattern so Phase 3's `stack_usage()` can measure
292    // the high-water mark (the deepest untouched 0xAA byte marks how far
293    // the task actually ran down its stack).
294    let base = stack.as_ptr() as usize;
295    let size = stack.len();
296    // The stack lives in the MPU-denied task-stack pool (plan.md §3.1);
297    // open a scratch window so the kernel can fill and initialize it
298    // before the task ever runs (no-op on arches without per-range
299    // guards). Held under a critical section so no other task can run in
300    // the window.
301    let sp = crate::critical::enter(|| {
302        crate::port::arch::scratch_open(base, size);
303        stack.fill(0xAA);
304        let sp =
305            crate::port::arch::init_task_stack(stack, entry as usize, arg as *const A as usize);
306        crate::port::arch::scratch_close();
307        sp
308    });
309    // plan.md Phase 17 (found via soak testing, same investigation as the
310    // `joiner`-reset fix in `tcb::register_full`): `register_full` calls
311    // `sched::ready_add`, making the new task immediately dispatchable —
312    // if it's higher priority than the spawning task, a tick can dispatch
313    // and run it to completion (`rivet_task_exit_core`) *before* this
314    // function ever reaches the `result_size`/`result_drop` stores below.
315    // `rivet_task_exit_core` then sees `result_size == 0`, so `size > 0`
316    // is false and the return value is silently never written (and a
317    // droppable `T` leaks, since `result_drop` is also still 0). Wrapping
318    // registration and the metadata stores in one critical section closes
319    // the window: nothing can dispatch the new task until both are
320    // published together.
321    let registered = crate::critical::enter(|| {
322        let id = tcb::register_full(sp, priority, base, size)?;
323        let t = tcb::get(id).expect("just registered");
324        let sz = core::mem::size_of::<T>();
325        debug_assert!(
326            sz <= 8,
327            "rivet: task return values > 8 bytes are not supported (got {sz})"
328        );
329        t.result_size.store(sz as u8, Ordering::Release);
330        t.result_drop.store(
331            if core::mem::needs_drop::<T>() {
332                // Cast through a function pointer: `fn` → usize (clippy
333                // fn_to_numeric_cast).
334                drop_in_place_erased::<T> as *const () as usize
335            } else {
336                0
337            },
338            Ordering::Release,
339        );
340        Some((id, t.generation.load(Ordering::Acquire)))
341    });
342    match registered {
343        Some((id, generation)) => {
344            #[cfg(feature = "trace")]
345            crate::trace::task_created(id as u16, priority, size as u32);
346            Ok(TaskHandle {
347                id: id as u16,
348                generation,
349            })
350        }
351        None => Err(SpawnError::RegistryFull),
352    }
353}
354
355/// Type-erased `drop_in_place` for a stored task result.
356fn drop_in_place_erased<T>(ptr: *mut u8) {
357    // SAFETY: the caller guarantees `ptr` points at a live `T`.
358    unsafe {
359        core::ptr::drop_in_place(ptr as *mut T);
360    }
361}
362
363/// Declare and spawn a preemptive task in one step.
364///
365/// ```ignore
366/// static CONFIG: MyConfig = MyConfig { ... };
367///
368/// rivet::spawn_ptask!(stack = 2048, priority = 3, entry = my_task, arg = CONFIG);
369///
370/// fn my_task(cfg: &'static MyConfig) -> ! {
371///     loop { /* runs with real preemption, own stack */ }
372/// }
373/// ```
374#[macro_export]
375macro_rules! spawn_ptask {
376    (stack = $stack_size:expr, priority = $prio:expr, entry = $entry:expr, arg = $arg:expr) => {{
377        // The stack comes from the kernel's task-stack pool (plan.md §3) so
378        // MPU/PMP guards can isolate it. On host builds (no pool) fall back
379        // to a per-invocation static.
380        $crate::preempt::spawn_ptask_impl($stack_size, $prio, $entry, &$arg, || {
381            static mut __RIVET_PTASK_STACK: $crate::preempt::Stack<$stack_size> =
382                $crate::preempt::Stack::new();
383            // SAFETY: this fallback stack is used exactly once, for the
384            // lifetime of the task (host builds only).
385            #[allow(static_mut_refs)]
386            unsafe {
387                &mut __RIVET_PTASK_STACK.0
388            }
389        })
390    }};
391}
392
393/// Measure a task stack's high-water mark (plan.md §2.7/Phase 3): stacks
394/// are filled with `0xAA` at spawn; the deepest byte a task wrote
395/// (anything else) marks how far it ran down. Returns the number of bytes
396/// used from the top.
397pub fn stack_usage(stack: &[u8]) -> usize {
398    let used = stack.iter().take_while(|&&b| b == 0xAA).count();
399    stack.len().saturating_sub(used)
400}
401
402/// Start the preemptive scheduler. Never returns — control transfers
403/// permanently to whichever task the scheduler selects first (and from
404/// there, forever between tasks via interrupt-driven context switches).
405///
406/// Call this from hart 0 only; secondary harts (plan.md Phase 19) use
407/// [`start_secondary_hart`] instead.
408///
409/// # Panics
410/// Panics if no preemptive tasks have been spawned.
411pub fn start() -> ! {
412    // Root cause (plan.md Phase 24), found on real dual-core hardware,
413    // **re-found in this exact form** (plan.md Phase 29) after the Phase
414    // 24 fix below turned out to still have a gap: the original fix
415    // wrapped only the `port::arch::critical_section` block below around
416    // the dispatch, leaving `crate::critical::enter`'s own interrupt mask
417    // (below) to release *before* that block's `irq_save` took effect —
418    // a real, if narrow, window between the two separate critical
419    // sections where this hart's local interrupts are genuinely enabled
420    // again. A tick or cross-hart IPI landing in that specific gap hits
421    // the identical failure the comment below describes. Closed by
422    // making the *entire* function body — the scheduling decision and
423    // the dispatch — one unbroken interrupt-masked region: entering
424    // `port::arch::critical_section` first, `crate::critical::enter`'s
425    // own nested local-mask-then-restore composes correctly inside it
426    // (`critical_section`'s docs: "Nested calls compose... its restore is
427    // a no-op, leaving the outermost call to actually re-enable"), so
428    // interrupts never actually re-enable until the dispatched task's own
429    // context takes over. Confirmed: `smp_latency_bench`'s forced-cross-
430    // core scenario (holder+waiter only, no other tasks) reproduced this
431    // panic on every attempt before this fix; see the QEMU/hardware
432    // re-verification this phase ran after the change.
433    crate::port::arch::critical_section(|| {
434        // plan.md Phase 19: the read-decide-commit sequence (pick a task,
435        // mark it Running, publish it as this hart's `CURRENT`) must be
436        // atomic across harts, not just locally-interrupt-safe — a
437        // secondary hart could be inside the identical sequence in
438        // `start_secondary_hart` concurrently. Single-hart boards get the
439        // same control flow with an always-uncontended CAS (see
440        // `critical.rs`'s module docs).
441        let first = crate::critical::enter(|| {
442            let first = sched::schedule().expect(
443                "rivet::preempt::start(): no preemptive tasks spawned (call rivet::init() \
444                 first, which spawns the async idle task, or spawn at least one via \
445                 spawn_ptask!)",
446            );
447            sched::set_current(first);
448            if let Some(t) = tcb::get(first) {
449                t.set_state(first, TaskState::Running);
450            }
451            // First dispatch: advance the RR start past this task (plan.md
452            // [B14]).
453            sched::on_dispatch(first);
454            first
455        });
456        // `sched::set_current(first)` (above) is visible the instant that
457        // inner `critical::enter` exits, but this hart's own CPU
458        // registers/stack aren't anywhere near `first`'s bootstrap state
459        // yet — a tick or IPI landing here, before interrupts are masked,
460        // would see `sched::current() == Some(first)` and, if it decides
461        // to switch, permanently rewrite `Tcb.sp` from a bootstrap marker
462        // to a real task id before `first` has ever run a single
463        // instruction. Because this whole function is now one continuous
464        // masked region (see the comment above `critical_section`), that
465        // window no longer exists — this comment describes what *would*
466        // happen without that framing, not a residual gap.
467        crate::exec_time::on_first_dispatch();
468        let first_tcb = tcb::get(first).unwrap();
469        // Enable memory protection for this task's stack — hart-local,
470        // done outside the cross-hart critical section like the rest of
471        // the arch dispatch.
472        crate::port::arch::on_switch_to(
473            first_tcb.stack_base.load(Ordering::Acquire),
474            first_tcb.stack_size.load(Ordering::Acquire),
475        );
476        let sp = first_tcb.sp.load(Ordering::Acquire);
477        // SAFETY: `sp` is the freshly-initialized first stack frame of
478        // the selected task (produced by `init_task_stack`);
479        // `start_first_task` consumes it exactly once and never returns.
480        // `critical_section`'s own irq-restore-on-return is consequently
481        // never reached — the dispatched task's arch-side bootstrap is
482        // responsible for re-enabling interrupts as part of actually
483        // starting to run, the same handoff `fresh_task_context`'s
484        // `Context.PS`/equivalent already does for the *ordinary*
485        // tick-driven first-dispatch path.
486        unsafe { crate::port::arch::start_first_task(sp) }
487    })
488}
489
490/// Start the preemptive scheduler on a **secondary** hart (plan.md
491/// Phase 19). Identical to [`start`] except a hart that finds nothing
492/// ready yet idles and retries instead of panicking: unlike hart 0 (which
493/// only starts after at least one task has been spawned), a secondary
494/// hart legitimately has nothing to do until some hart makes a task ready
495/// and IPIs it via `port::arch::request_reschedule_on`.
496///
497/// Never returns. No-op-forever (idles) on a single-hart board, since
498/// nothing ever calls it there — `rivet-rt`'s hart bring-up only invokes
499/// this for harts `1..RIVET_MAX_HARTS`.
500pub fn start_secondary_hart() -> ! {
501    // Same fix, same reason as `start`'s own (plan.md Phase 24, re-closed
502    // Phase 29 — see `start`'s comment for the full story): the
503    // scheduling decision and the dispatch must be one unbroken
504    // interrupt-masked region on *this* hart, not two separate critical
505    // sections with a real gap between them. `idle()` still needs to run
506    // with interrupts genuinely enabled (it's a wait-for-interrupt, it
507    // would never wake otherwise) — masked only starts once a candidate
508    // is actually found, per loop iteration, not around the whole loop.
509    loop {
510        let dispatched = crate::port::arch::critical_section(|| {
511            let first = crate::critical::enter(|| {
512                let first = sched::schedule()?;
513                sched::set_current(first);
514                if let Some(t) = tcb::get(first) {
515                    t.set_state(first, TaskState::Running);
516                }
517                sched::on_dispatch(first);
518                Some(first)
519            });
520            let Some(first) = first else {
521                return false;
522            };
523            let first_tcb = tcb::get(first).unwrap();
524            // No `exec_time::on_first_dispatch()` call here: hart 0's
525            // `start()` already stamped `BOOT_CYCLE`/`WALLCLOCK_BOOT_US`
526            // once for the whole system (plan.md Phase 19 §6 — exec-time
527            // accounting stays a single shared boot epoch, not per-hart;
528            // calling it again here would rewind `LAST_DISPATCH` and skew
529            // every other task's busy-cycle accounting).
530            crate::port::arch::on_switch_to(
531                first_tcb.stack_base.load(Ordering::Acquire),
532                first_tcb.stack_size.load(Ordering::Acquire),
533            );
534            let sp = first_tcb.sp.load(Ordering::Acquire);
535            // SAFETY: `sp` is the freshly-initialized first stack frame
536            // of the selected task (produced by `init_task_stack`);
537            // `start_first_task` consumes it exactly once and never
538            // returns. `critical_section`'s own irq-restore-on-return is
539            // consequently never reached on this path — the dispatched
540            // task's arch-side bootstrap re-enables interrupts itself,
541            // same as `start`'s identical call.
542            unsafe { crate::port::arch::start_first_task(sp) }
543        });
544        if !dispatched {
545            crate::port::arch::idle();
546        }
547    }
548}
549
550/// Permanently remove the current preemptive task from scheduling. Useful
551/// for a task that does bounded work and then has nothing left to do —
552/// parking (rather than spinning forever at its original priority) lets
553/// lower-priority tasks actually run.
554///
555/// # Panics
556/// Panics if called outside of a preemptive task context.
557/// Block the current preemptive task for `ms` milliseconds (plan.md §5.6):
558/// registers a deadline in the per-task queue, blocks, and lets the timer
559/// tick wake it. No-op outside a preemptive task.
560pub fn sleep_ms(ms: u64) {
561    let deadline = crate::port::board::now_us().wrapping_add(ms.saturating_mul(1000));
562    sleep_until(deadline);
563}
564
565/// Block the calling preemptive task until the absolute time `deadline_us`
566/// (plan.md §5.6 / Phase 11). `sleep_ms` is `sleep_until(now + ms*1000)`;
567/// [`crate::deadlines::wait_period`] uses this directly with a
568/// drift-corrected deadline so periodic jitter doesn't accumulate. No-op
569/// outside a preemptive task context. If `deadline_us` has already
570/// passed, still yields once (bounded, not a busy spin) rather than
571/// returning immediately.
572pub fn sleep_until(deadline_us: u64) {
573    let Some(me) = sched::current() else {
574        return;
575    };
576    // `block_current()` and `register_ptask_deadline()` must commit
577    // together: a tick landing between them would see the task already
578    // Blocked but with no deadline registered yet, switch away from it
579    // (correctly — it's not ready), and then never find anything to wake
580    // it later — a permanent lost-wakeup hang. Same shape, same fix as
581    // the mutex slow path (`mutex.rs`'s own `critical::enter` around the
582    // equivalent boost/add_waiter/register/block sequence) and the
583    // lifecycle join path; this call site was the one left over from
584    // before that pattern was established.
585    crate::critical::enter(|| {
586        sched::block_current();
587        let _ = crate::timer::register_ptask_deadline(deadline_us, me);
588    });
589    crate::port::arch::request_reschedule();
590    crate::timer::cancel_ptask_deadline(me);
591}
592
593pub fn park_forever() -> ! {
594    sched::current().expect("park_forever() outside preemptive task context");
595    sched::block_current();
596    loop {
597        crate::port::arch::request_reschedule();
598    }
599}
600
601/// Called from the arch trap/exception handler (timer tick, or a software
602/// interrupt triggered by [`crate::port::arch::request_reschedule`]) with the interrupted
603/// task's just-saved stack pointer. Consults the scheduler and returns the
604/// stack pointer the arch layer should actually resume — either the same
605/// one (no reschedule needed) or a different task's (real preemption /
606/// voluntary switch).
607///
608/// If the preemptive tier hasn't started yet ([`start`] not called),
609/// returns `interrupted_sp` unchanged.
610pub fn on_tick(interrupted_sp: usize) -> usize {
611    #[cfg(feature = "latency-histograms")]
612    let __latency_start = crate::port::arch::cycle_count();
613    let result = on_tick_impl(interrupted_sp);
614    #[cfg(feature = "latency-histograms")]
615    crate::latency::record(
616        crate::latency::Kind::DispatchDecision,
617        crate::port::arch::cycle_count().wrapping_sub(__latency_start),
618    );
619    // Real bug, found via the debugger it was corrupting the timing of:
620    // `on_tick_locked` used to call `crate::trace::context_switch(...)`
621    // directly, inside `on_tick_impl`'s own `critical::enter` (PRIMASK
622    // masked, i.e. the *whole system* stopped, not just this hart's local
623    // interrupts). A trace emission is a blocking, byte-at-a-time polling
624    // UART write — tens of microseconds per byte, over a millisecond for
625    // a whole frame at 115200 baud — which is worse than the entire 1kHz
626    // tick period. Two equal-priority tasks that round-robin every tick
627    // (`should_preempt`'s documented "equal priority round-robins on
628    // tick" semantics) hit this on *every single dispatch*: the outgoing
629    // task's ContextSwitch trace_write alone consumed the whole tick
630    // budget, so the incoming task got dispatched, immediately had its
631    // one tick's worth of runway eaten by the *next* tick's own blocking
632    // trace_write (already pending the instant PRIMASK lifted), and never
633    // advanced past its own entry point — confirmed via a live GDB
634    // capture showing two worker tasks permanently stuck with their saved
635    // PC at their function's first instruction, cumulative spin/sleep
636    // completions stuck at 0 after several real seconds of uptime. Same
637    // hazard class as the reannounce comment below (`docs/wcet.md` §6.1);
638    // this is the sequel finding it in the one spot that had already
639    // shipped it inside the lock instead of just being tempted to.
640    // `PENDING_CTX_SWITCH` carries the (prev,next) pair out of the locked
641    // region so the actual UART write happens here, after PRIMASK lifts.
642    #[cfg(feature = "trace")]
643    {
644        use core::sync::atomic::Ordering;
645        let packed = PENDING_CTX_SWITCH.swap(u32::MAX, Ordering::Relaxed);
646        if packed != u32::MAX {
647            let prev = (packed >> 16) as u16;
648            let next = (packed & 0xFFFF) as u16;
649            crate::trace::context_switch(prev, next, crate::trace::SwitchReason::Preempted);
650        }
651    }
652    // Deliberately outside on_tick_impl's own critical::enter: a full
653    // task-table scan plus one trace_write per live task is far more
654    // than a critical section should ever hold (this exact class of
655    // "console/trace I/O under critical::enter" cost is a real, measured
656    // hazard — see docs/wcet.md §6.1 in this workspace's own real-
657    // hardware WCET analysis). Gated to roughly once every two seconds
658    // at the default 1kHz tick, not every tick.
659    #[cfg(feature = "trace")]
660    {
661        use core::sync::atomic::{AtomicU32, Ordering};
662        static REANNOUNCE_TICK: AtomicU32 = AtomicU32::new(0);
663        if REANNOUNCE_TICK.fetch_add(1, Ordering::Relaxed).is_multiple_of(2000) {
664            crate::trace::reannounce_all_tasks();
665            crate::trace::reannounce_stream_header();
666        }
667    }
668    result
669}
670
671/// Packs `(prev_task << 16) | next_task` for the most recent tick-driven
672/// switch, `u32::MAX` = none pending. Written (Relaxed — single-hart-tick-
673/// owner, same reasoning as `REANNOUNCE_TICK`) from inside
674/// `on_tick_locked`'s critical section, read/cleared from [`on_tick`]
675/// after that section has already released PRIMASK — see `on_tick`'s own
676/// doc for why this indirection exists at all.
677#[cfg(feature = "trace")]
678static PENDING_CTX_SWITCH: core::sync::atomic::AtomicU32 = core::sync::atomic::AtomicU32::new(u32::MAX);
679
680/// plan.md Phase 19: the entire read-decide-commit sequence below (read
681/// `sched::current()`, decide via `schedule()`/`should_preempt`, commit via
682/// `set_state`/`set_current`/`on_dispatch`) must be atomic across harts —
683/// two harts ticking concurrently could otherwise both claim the same
684/// ready task. Wrapping it in `critical::enter` closes that gap; on a
685/// single-hart board the wrap is a no-op (see `critical.rs`'s module
686/// docs), so this is behavior-preserving there.
687fn on_tick_impl(interrupted_sp: usize) -> usize {
688    crate::critical::enter(|| on_tick_locked(interrupted_sp))
689}
690
691fn on_tick_locked(interrupted_sp: usize) -> usize {
692    let Some(running) = sched::current() else {
693        return interrupted_sp;
694    };
695
696    if let Some(t) = tcb::get(running) {
697        t.sp.store(interrupted_sp, Ordering::Release);
698    }
699
700    let Some(candidate) = sched::schedule() else {
701        return interrupted_sp;
702    };
703
704    // Watermark overflow check (plan.md §3.3): the outgoing task's lowest
705    // stack word must still be the 0xAA fill pattern. Catches overflow
706    // that the MPU/PMP guards miss (RISC-V tasks beyond the PMP budget, or
707    // a large stack array jumping over a small guard band).
708    if let Some(t) = tcb::get(running) {
709        let base = t.stack_base.load(Ordering::Acquire);
710        let size = t.stack_size.load(Ordering::Acquire);
711        if base != 0 && size >= 4 {
712            // SAFETY: reading the running task's own stack is always
713            // allowed (it is the MPU-enabled current stack).
714            let lowest = unsafe { core::ptr::read_volatile(base as *const u32) };
715            if lowest != 0xAAAA_AAAA {
716                let info = crate::fault::FaultInfo {
717                    task_id: Some(running),
718                    kind: crate::fault::FaultKind::StackOverflow,
719                    address: base,
720                    pc: 0,
721                };
722                return crate::fault::on_fault(&info);
723            }
724        }
725    }
726
727    // CPU-budget check (plan.md Phase 11): only meaningful for a task
728    // that's still actually Running (a task that just blocked itself, via
729    // `wait_period`'s own `sleep_until`, resets its budget window on the
730    // *next* period start, not here).
731    if tcb::get(running).map(|t| t.state()) == Some(TaskState::Running)
732        && crate::deadlines::check_budget(running)
733    {
734        let info = crate::fault::FaultInfo {
735            task_id: Some(running),
736            kind: crate::fault::FaultKind::BudgetExceeded,
737            address: 0,
738            pc: 0,
739        };
740        return crate::fault::on_fault(&info);
741    }
742
743    // A task that just blocked itself (e.g. park_forever(), or a
744    // PriorityMutex wait) must be switched away from unconditionally —
745    // should_preempt()'s priority comparison only makes sense between two
746    // tasks that could both legitimately keep running. A Blocked task
747    // can't "keep running" at all, regardless of whether the candidate's
748    // priority is lower (e.g. falling through to the priority-0 async
749    // idle task). Without this check, once every task at the blocked
750    // task's priority level is also blocked, on_tick keeps returning
751    // interrupted_sp forever — spinning inside the blocked task's own
752    // park loop instead of ever handing off to the (lower-priority, but
753    // only-ready) candidate.
754    let running_blocked = tcb::get(running)
755        .map(|t| t.state() == TaskState::Blocked)
756        .unwrap_or(true);
757
758    if !running_blocked && !sched::should_preempt(candidate, running) {
759        return interrupted_sp;
760    }
761    if running_blocked && candidate == running {
762        // Nothing else is ready; stay parked (spurious wake or no other work).
763        return interrupted_sp;
764    }
765
766    if let Some(t) = tcb::get(running) {
767        if t.state() == TaskState::Running {
768            t.set_state(running, TaskState::Ready);
769        }
770    }
771    crate::exec_time::on_switch(running);
772    let to_tcb = tcb::get(candidate).unwrap();
773    to_tcb.set_state(candidate, TaskState::Running);
774    // NOT a direct `crate::trace::context_switch(...)` call here — see
775    // `on_tick`'s doc comment on `PENDING_CTX_SWITCH` for why a blocking
776    // UART write from inside this critical section is a real, confirmed
777    // starvation bug, not just a theoretical latency concern.
778    #[cfg(feature = "trace")]
779    PENDING_CTX_SWITCH.store(
780        ((running as u32) << 16) | (candidate as u32),
781        core::sync::atomic::Ordering::Relaxed,
782    );
783    sched::set_current(candidate);
784    // An actual switch occurred: advance the RR start past the dispatched
785    // task (plan.md [B14] — never advance on no-switch ticks), and enable
786    // memory protection for the newly-running task's stack (plan.md §3.1).
787    sched::on_dispatch(candidate);
788    crate::port::arch::on_switch_to(
789        to_tcb.stack_base.load(Ordering::Acquire),
790        to_tcb.stack_size.load(Ordering::Acquire),
791    );
792    to_tcb.sp.load(Ordering::Acquire)
793}
794
795#[cfg(test)]
796mod stack_tests {
797    use super::*;
798
799    #[test]
800    fn stack_usage_measures_fill_pattern() {
801        let mut stack = [0xAAu8; 512];
802        assert_eq!(stack_usage(&stack), 0, "untouched stack uses nothing");
803        // The stack grows DOWN from the top: a task that ran 256 bytes
804        // deep leaves the bottom 256 bytes as untouched 0xAA.
805        stack[256..].fill(0x00);
806        assert_eq!(stack_usage(&stack), 256);
807    }
808
809    #[test]
810    #[should_panic(expected = "task stack too small")]
811    fn spawn_rejects_too_small_stack() {
812        crate::kernel_test! {
813            static mut TINY: [u8; 32] = [0; 32];
814            fn entry(_: &'static ()) -> ! { loop { crate::port::arch::request_reschedule(); } }
815            static UNIT: () = ();
816            // SAFETY: TINY is only used here, before the scheduler runs.
817            unsafe {
818                // SAFETY: TINY is only used here, before the scheduler
819                // runs; addr_of_mut! avoids a reference to the static.
820                let stack = &mut (*core::ptr::addr_of_mut!(TINY));
821                let _ = spawn(stack, 1, entry, &UNIT);
822            }
823        }
824    }
825}