pub mod lifecycle;
pub mod mutex;
pub mod sched;
pub mod stack_pool;
pub mod tcb;
pub use mutex::{PriorityMutex, PriorityMutexGuard};
pub use tcb::TaskState;
use crate::sync::atomic::Ordering;
#[repr(C, align(16))]
pub struct Stack<const SIZE: usize>(pub [u8; SIZE]);
impl<const SIZE: usize> Stack<SIZE> {
pub const fn new() -> Self {
Self([0; SIZE])
}
}
impl<const SIZE: usize> Default for Stack<SIZE> {
fn default() -> Self {
Self::new()
}
}
#[doc(hidden)]
pub fn spawn_ptask_impl<T: 'static + Send, A: 'static, F: Fn() -> &'static mut [u8]>(
stack_size: usize,
priority: u8,
entry: fn(&'static A) -> T,
arg: &'static A,
#[allow(unused_variables)] fallback: F,
) -> Result<TaskHandle, SpawnError> {
let stack = match crate::preempt::stack_pool::alloc_stack(stack_size) {
Some(s) => s,
None => {
#[cfg(not(feature = "host-port"))]
return Err(SpawnError::StackPoolFull);
#[cfg(feature = "host-port")]
fallback()
}
};
unsafe { spawn(stack, priority, entry, arg) }
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct TaskHandle {
pub id: u16,
pub generation: u32,
}
pub use lifecycle::JoinError;
impl TaskHandle {
pub fn is_valid(&self) -> bool {
tcb::get(self.id as usize)
.map(|t| t.generation.load(Ordering::Acquire) == self.generation)
.unwrap_or(false)
}
pub fn request_stop(&self) -> bool {
match tcb::get(self.id as usize) {
Some(t) if t.generation.load(Ordering::Acquire) == self.generation => {
t.stop_requested.store(true, Ordering::Release);
true
}
_ => false,
}
}
pub fn join<T: 'static + Send>(&self) -> Result<T, JoinError> {
lifecycle::join_task::<T>(self)
}
pub fn set_period_us(&self, period_us: u32) {
if self.is_valid() {
crate::deadlines::set_period_us(self.id as usize, period_us);
}
}
pub fn set_budget_us(&self, budget_us: u32) {
if self.is_valid() {
crate::deadlines::set_budget_us(self.id as usize, budget_us);
}
}
pub fn despawn(&self) -> bool {
let Some(t) = tcb::get(self.id as usize) else {
return false;
};
if t.generation.load(Ordering::Acquire) != self.generation {
return false;
}
if sched::current() == Some(self.id as usize) {
return false; }
if !t.used.load(Ordering::Acquire) {
return false;
}
crate::critical::enter(|| crate::preempt::sched::ready_remove(self.id as usize));
let drop_fn = t.result_drop.load(Ordering::Acquire);
if drop_fn != 0 {
let f: fn(*mut u8) = unsafe { core::mem::transmute(drop_fn) };
f(t.result_buf.get() as *mut u8);
t.result_drop.store(0, Ordering::Release);
}
t.state.store(tcb::READY, Ordering::Release);
t.exited.store(false, Ordering::Release);
t.result_size.store(0, Ordering::Release);
t.stop_requested.store(false, Ordering::Release);
t.used.store(false, Ordering::Release);
let stack = t.stack_info();
if let Some((base, size)) = stack {
if base != 0 && size != 0 {
let slice: &'static mut [u8] =
unsafe { core::slice::from_raw_parts_mut(base as *mut u8, size) };
crate::preempt::stack_pool::release_stack(slice);
}
}
t.stack_base.store(0, Ordering::Release);
t.stack_size.store(0, Ordering::Release);
true
}
pub fn pause(&self) -> bool {
let Some(t) = tcb::get(self.id as usize) else {
return false;
};
if t.generation.load(Ordering::Acquire) != self.generation {
return false;
}
crate::critical::enter(|| {
let was_ready = t
.state
.compare_exchange(
tcb::READY,
tcb::SUSPENDED,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok();
if was_ready {
crate::preempt::sched::ready_remove(self.id as usize);
}
was_ready
})
}
pub fn resume(&self) -> bool {
let Some(t) = tcb::get(self.id as usize) else {
return false;
};
if t.generation.load(Ordering::Acquire) != self.generation {
return false;
}
crate::critical::enter(|| {
let was_suspended = t
.state
.compare_exchange(
tcb::SUSPENDED,
tcb::READY,
Ordering::AcqRel,
Ordering::Acquire,
)
.is_ok();
if was_suspended {
crate::preempt::sched::ready_add(self.id as usize);
}
was_suspended
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SpawnError {
RegistryFull,
StackPoolFull,
}
pub unsafe fn spawn<T: 'static + Send, A: 'static>(
stack: &'static mut [u8],
priority: u8,
entry: fn(&'static A) -> T,
arg: &'static A,
) -> Result<TaskHandle, SpawnError> {
assert!(
stack.len() >= crate::port::arch::min_task_stack(),
"rivet: task stack too small: {} bytes < arch minimum {} (context-switch frame + entry trampoline; plan.md §2.7)",
stack.len(),
crate::port::arch::min_task_stack()
);
let base = stack.as_ptr() as usize;
let size = stack.len();
let sp = crate::critical::enter(|| {
crate::port::arch::scratch_open(base, size);
stack.fill(0xAA);
let sp =
crate::port::arch::init_task_stack(stack, entry as usize, arg as *const A as usize);
crate::port::arch::scratch_close();
sp
});
let registered = crate::critical::enter(|| {
let id = tcb::register_full(sp, priority, base, size)?;
let t = tcb::get(id).expect("just registered");
let sz = core::mem::size_of::<T>();
debug_assert!(
sz <= 8,
"rivet: task return values > 8 bytes are not supported (got {sz})"
);
t.result_size.store(sz as u8, Ordering::Release);
t.result_drop.store(
if core::mem::needs_drop::<T>() {
drop_in_place_erased::<T> as *const () as usize
} else {
0
},
Ordering::Release,
);
Some((id, t.generation.load(Ordering::Acquire)))
});
match registered {
Some((id, generation)) => {
#[cfg(feature = "trace")]
crate::trace::task_created(id as u16, priority, size as u32);
Ok(TaskHandle {
id: id as u16,
generation,
})
}
None => Err(SpawnError::RegistryFull),
}
}
fn drop_in_place_erased<T>(ptr: *mut u8) {
unsafe {
core::ptr::drop_in_place(ptr as *mut T);
}
}
#[macro_export]
macro_rules! spawn_ptask {
(stack = $stack_size:expr, priority = $prio:expr, entry = $entry:expr, arg = $arg:expr) => {{
$crate::preempt::spawn_ptask_impl($stack_size, $prio, $entry, &$arg, || {
static mut __RIVET_PTASK_STACK: $crate::preempt::Stack<$stack_size> =
$crate::preempt::Stack::new();
#[allow(static_mut_refs)]
unsafe {
&mut __RIVET_PTASK_STACK.0
}
})
}};
}
pub fn stack_usage(stack: &[u8]) -> usize {
let used = stack.iter().take_while(|&&b| b == 0xAA).count();
stack.len().saturating_sub(used)
}
pub fn start() -> ! {
crate::port::arch::critical_section(|| {
let first = crate::critical::enter(|| {
let first = sched::schedule().expect(
"rivet::preempt::start(): no preemptive tasks spawned (call rivet::init() \
first, which spawns the async idle task, or spawn at least one via \
spawn_ptask!)",
);
sched::set_current(first);
if let Some(t) = tcb::get(first) {
t.set_state(first, TaskState::Running);
}
sched::on_dispatch(first);
first
});
crate::exec_time::on_first_dispatch();
let first_tcb = tcb::get(first).unwrap();
crate::port::arch::on_switch_to(
first_tcb.stack_base.load(Ordering::Acquire),
first_tcb.stack_size.load(Ordering::Acquire),
);
let sp = first_tcb.sp.load(Ordering::Acquire);
unsafe { crate::port::arch::start_first_task(sp) }
})
}
pub fn start_secondary_hart() -> ! {
loop {
let dispatched = crate::port::arch::critical_section(|| {
let first = crate::critical::enter(|| {
let first = sched::schedule()?;
sched::set_current(first);
if let Some(t) = tcb::get(first) {
t.set_state(first, TaskState::Running);
}
sched::on_dispatch(first);
Some(first)
});
let Some(first) = first else {
return false;
};
let first_tcb = tcb::get(first).unwrap();
crate::port::arch::on_switch_to(
first_tcb.stack_base.load(Ordering::Acquire),
first_tcb.stack_size.load(Ordering::Acquire),
);
let sp = first_tcb.sp.load(Ordering::Acquire);
unsafe { crate::port::arch::start_first_task(sp) }
});
if !dispatched {
crate::port::arch::idle();
}
}
}
pub fn sleep_ms(ms: u64) {
let deadline = crate::port::board::now_us().wrapping_add(ms.saturating_mul(1000));
sleep_until(deadline);
}
pub fn sleep_until(deadline_us: u64) {
let Some(me) = sched::current() else {
return;
};
crate::critical::enter(|| {
sched::block_current();
let _ = crate::timer::register_ptask_deadline(deadline_us, me);
});
crate::port::arch::request_reschedule();
crate::timer::cancel_ptask_deadline(me);
}
pub fn park_forever() -> ! {
sched::current().expect("park_forever() outside preemptive task context");
sched::block_current();
loop {
crate::port::arch::request_reschedule();
}
}
pub fn on_tick(interrupted_sp: usize) -> usize {
#[cfg(feature = "latency-histograms")]
let __latency_start = crate::port::arch::cycle_count();
let result = on_tick_impl(interrupted_sp);
#[cfg(feature = "latency-histograms")]
crate::latency::record(
crate::latency::Kind::DispatchDecision,
crate::port::arch::cycle_count().wrapping_sub(__latency_start),
);
#[cfg(feature = "trace")]
{
use core::sync::atomic::Ordering;
let packed = PENDING_CTX_SWITCH.swap(u32::MAX, Ordering::Relaxed);
if packed != u32::MAX {
let prev = (packed >> 16) as u16;
let next = (packed & 0xFFFF) as u16;
crate::trace::context_switch(prev, next, crate::trace::SwitchReason::Preempted);
}
}
#[cfg(feature = "trace")]
{
use core::sync::atomic::{AtomicU32, Ordering};
static REANNOUNCE_TICK: AtomicU32 = AtomicU32::new(0);
if REANNOUNCE_TICK.fetch_add(1, Ordering::Relaxed).is_multiple_of(2000) {
crate::trace::reannounce_all_tasks();
crate::trace::reannounce_stream_header();
}
}
result
}
#[cfg(feature = "trace")]
static PENDING_CTX_SWITCH: core::sync::atomic::AtomicU32 = core::sync::atomic::AtomicU32::new(u32::MAX);
fn on_tick_impl(interrupted_sp: usize) -> usize {
crate::critical::enter(|| on_tick_locked(interrupted_sp))
}
fn on_tick_locked(interrupted_sp: usize) -> usize {
let Some(running) = sched::current() else {
return interrupted_sp;
};
if let Some(t) = tcb::get(running) {
t.sp.store(interrupted_sp, Ordering::Release);
}
let Some(candidate) = sched::schedule() else {
return interrupted_sp;
};
if let Some(t) = tcb::get(running) {
let base = t.stack_base.load(Ordering::Acquire);
let size = t.stack_size.load(Ordering::Acquire);
if base != 0 && size >= 4 {
let lowest = unsafe { core::ptr::read_volatile(base as *const u32) };
if lowest != 0xAAAA_AAAA {
let info = crate::fault::FaultInfo {
task_id: Some(running),
kind: crate::fault::FaultKind::StackOverflow,
address: base,
pc: 0,
};
return crate::fault::on_fault(&info);
}
}
}
if tcb::get(running).map(|t| t.state()) == Some(TaskState::Running)
&& crate::deadlines::check_budget(running)
{
let info = crate::fault::FaultInfo {
task_id: Some(running),
kind: crate::fault::FaultKind::BudgetExceeded,
address: 0,
pc: 0,
};
return crate::fault::on_fault(&info);
}
let running_blocked = tcb::get(running)
.map(|t| t.state() == TaskState::Blocked)
.unwrap_or(true);
if !running_blocked && !sched::should_preempt(candidate, running) {
return interrupted_sp;
}
if running_blocked && candidate == running {
return interrupted_sp;
}
if let Some(t) = tcb::get(running) {
if t.state() == TaskState::Running {
t.set_state(running, TaskState::Ready);
}
}
crate::exec_time::on_switch(running);
let to_tcb = tcb::get(candidate).unwrap();
to_tcb.set_state(candidate, TaskState::Running);
#[cfg(feature = "trace")]
PENDING_CTX_SWITCH.store(
((running as u32) << 16) | (candidate as u32),
core::sync::atomic::Ordering::Relaxed,
);
sched::set_current(candidate);
sched::on_dispatch(candidate);
crate::port::arch::on_switch_to(
to_tcb.stack_base.load(Ordering::Acquire),
to_tcb.stack_size.load(Ordering::Acquire),
);
to_tcb.sp.load(Ordering::Acquire)
}
#[cfg(test)]
mod stack_tests {
use super::*;
#[test]
fn stack_usage_measures_fill_pattern() {
let mut stack = [0xAAu8; 512];
assert_eq!(stack_usage(&stack), 0, "untouched stack uses nothing");
stack[256..].fill(0x00);
assert_eq!(stack_usage(&stack), 256);
}
#[test]
#[should_panic(expected = "task stack too small")]
fn spawn_rejects_too_small_stack() {
crate::kernel_test! {
static mut TINY: [u8; 32] = [0; 32];
fn entry(_: &'static ()) -> ! { loop { crate::port::arch::request_reschedule(); } }
static UNIT: () = ();
unsafe {
let stack = &mut (*core::ptr::addr_of_mut!(TINY));
let _ = spawn(stack, 1, entry, &UNIT);
}
}
}
}