rivet/
waker.rs

1//! Zero-allocation waker using atomic priority bitmaps.
2//!
3//! Each task has a `(priority, index_in_priority)` pair. When a waker fires,
4//! it atomically sets the corresponding bit in the per-priority queue and
5//! the global ready bitmap. The executor finds the next task in O(1) using
6//! `leading_zeros` on the bitmap.
7
8// FQN atomics used below.
9use core::task::{RawWaker, RawWakerVTable, Waker};
10
11/// Bit i set => priority level i has at least one ready task.
12///
13/// Under `--cfg loom` these globals live in `loom::lazy_static!` (loom's
14/// atomics are not const-constructible); loom resets them between models.
15#[cfg(not(loom))]
16static READY_BITMAP: crate::sync::atomic::AtomicU32 = crate::sync::atomic::AtomicU32::new(0);
17#[cfg(loom)]
18loom::lazy_static! {
19    static ref READY_BITMAP: crate::sync::atomic::AtomicU32 = crate::sync::atomic::AtomicU32::new(0);
20}
21
22/// Per-priority ready task bitmasks.
23#[cfg(not(loom))]
24static PRIORITY_QUEUES: [crate::sync::atomic::AtomicU32; 32] = [
25    crate::sync::atomic::AtomicU32::new(0),
26    crate::sync::atomic::AtomicU32::new(0),
27    crate::sync::atomic::AtomicU32::new(0),
28    crate::sync::atomic::AtomicU32::new(0),
29    crate::sync::atomic::AtomicU32::new(0),
30    crate::sync::atomic::AtomicU32::new(0),
31    crate::sync::atomic::AtomicU32::new(0),
32    crate::sync::atomic::AtomicU32::new(0),
33    crate::sync::atomic::AtomicU32::new(0),
34    crate::sync::atomic::AtomicU32::new(0),
35    crate::sync::atomic::AtomicU32::new(0),
36    crate::sync::atomic::AtomicU32::new(0),
37    crate::sync::atomic::AtomicU32::new(0),
38    crate::sync::atomic::AtomicU32::new(0),
39    crate::sync::atomic::AtomicU32::new(0),
40    crate::sync::atomic::AtomicU32::new(0),
41    crate::sync::atomic::AtomicU32::new(0),
42    crate::sync::atomic::AtomicU32::new(0),
43    crate::sync::atomic::AtomicU32::new(0),
44    crate::sync::atomic::AtomicU32::new(0),
45    crate::sync::atomic::AtomicU32::new(0),
46    crate::sync::atomic::AtomicU32::new(0),
47    crate::sync::atomic::AtomicU32::new(0),
48    crate::sync::atomic::AtomicU32::new(0),
49    crate::sync::atomic::AtomicU32::new(0),
50    crate::sync::atomic::AtomicU32::new(0),
51    crate::sync::atomic::AtomicU32::new(0),
52    crate::sync::atomic::AtomicU32::new(0),
53    crate::sync::atomic::AtomicU32::new(0),
54    crate::sync::atomic::AtomicU32::new(0),
55    crate::sync::atomic::AtomicU32::new(0),
56    crate::sync::atomic::AtomicU32::new(0),
57];
58#[cfg(loom)]
59loom::lazy_static! {
60    static ref PRIORITY_QUEUES: [crate::sync::atomic::AtomicU32; 32] = [
61    crate::sync::atomic::AtomicU32::new(0),
62    crate::sync::atomic::AtomicU32::new(0),
63    crate::sync::atomic::AtomicU32::new(0),
64    crate::sync::atomic::AtomicU32::new(0),
65    crate::sync::atomic::AtomicU32::new(0),
66    crate::sync::atomic::AtomicU32::new(0),
67    crate::sync::atomic::AtomicU32::new(0),
68    crate::sync::atomic::AtomicU32::new(0),
69    crate::sync::atomic::AtomicU32::new(0),
70    crate::sync::atomic::AtomicU32::new(0),
71    crate::sync::atomic::AtomicU32::new(0),
72    crate::sync::atomic::AtomicU32::new(0),
73    crate::sync::atomic::AtomicU32::new(0),
74    crate::sync::atomic::AtomicU32::new(0),
75    crate::sync::atomic::AtomicU32::new(0),
76    crate::sync::atomic::AtomicU32::new(0),
77    crate::sync::atomic::AtomicU32::new(0),
78    crate::sync::atomic::AtomicU32::new(0),
79    crate::sync::atomic::AtomicU32::new(0),
80    crate::sync::atomic::AtomicU32::new(0),
81    crate::sync::atomic::AtomicU32::new(0),
82    crate::sync::atomic::AtomicU32::new(0),
83    crate::sync::atomic::AtomicU32::new(0),
84    crate::sync::atomic::AtomicU32::new(0),
85    crate::sync::atomic::AtomicU32::new(0),
86    crate::sync::atomic::AtomicU32::new(0),
87    crate::sync::atomic::AtomicU32::new(0),
88    crate::sync::atomic::AtomicU32::new(0),
89    crate::sync::atomic::AtomicU32::new(0),
90    crate::sync::atomic::AtomicU32::new(0),
91    crate::sync::atomic::AtomicU32::new(0),
92    crate::sync::atomic::AtomicU32::new(0),
93    ];
94}
95
96/// CPU flag set when waker fires from ISR context.
97#[cfg(not(loom))]
98pub(crate) static EXECUTOR_PEND_FLAG: crate::sync::atomic::AtomicU32 =
99    crate::sync::atomic::AtomicU32::new(0);
100#[cfg(loom)]
101loom::lazy_static! {
102    pub(crate) static ref EXECUTOR_PEND_FLAG: crate::sync::atomic::AtomicU32 = crate::sync::atomic::AtomicU32::new(0);
103}
104
105/// Reset the waker state (for testing).
106pub fn reset() {
107    READY_BITMAP.store(0, crate::sync::atomic::Ordering::Release);
108    for q in PRIORITY_QUEUES.iter() {
109        q.store(0, crate::sync::atomic::Ordering::Release);
110    }
111    EXECUTOR_PEND_FLAG.store(0, crate::sync::atomic::Ordering::Release);
112}
113
114/// Mark a task as ready.
115#[doc(hidden)]
116pub fn mark_ready(id: crate::task::TaskId) {
117    let mask = 1u32 << (id.index() & 0x1F);
118    PRIORITY_QUEUES[id.priority() as usize].fetch_or(mask, crate::sync::atomic::Ordering::Release);
119    READY_BITMAP.fetch_or(
120        1u32 << id.priority(),
121        crate::sync::atomic::Ordering::Release,
122    );
123    EXECUTOR_PEND_FLAG.store(1, crate::sync::atomic::Ordering::Release);
124}
125
126/// Check if any tasks are pending (used before sleep to avoid race).
127#[doc(hidden)]
128pub fn has_pending() -> bool {
129    EXECUTOR_PEND_FLAG.load(crate::sync::atomic::Ordering::Acquire) != 0
130        || READY_BITMAP.load(crate::sync::atomic::Ordering::Acquire) != 0
131}
132
133/// Clear the pend flag. Called when executor starts a new polling round.
134#[doc(hidden)]
135pub fn clear_pend() {
136    EXECUTOR_PEND_FLAG.store(0, crate::sync::atomic::Ordering::Release);
137}
138
139/// Dequeue the highest-priority ready task. Returns its [`TaskId`] or
140/// `None`.
141#[doc(hidden)]
142pub fn next_ready() -> Option<crate::task::TaskId> {
143    let bitmap = READY_BITMAP.load(crate::sync::atomic::Ordering::Acquire);
144    if bitmap == 0 {
145        return None;
146    }
147
148    let prio = (31 - bitmap.leading_zeros()) as u8;
149    let queue = PRIORITY_QUEUES[prio as usize].load(crate::sync::atomic::Ordering::Acquire);
150
151    if queue == 0 {
152        READY_BITMAP.fetch_and(!(1u32 << prio), crate::sync::atomic::Ordering::AcqRel);
153        return None;
154    }
155
156    let bit = queue & queue.wrapping_neg();
157    let index = bit.trailing_zeros() as u8;
158
159    let prev =
160        PRIORITY_QUEUES[prio as usize].fetch_and(!bit, crate::sync::atomic::Ordering::AcqRel);
161
162    if prev == bit {
163        READY_BITMAP.fetch_and(!(1u32 << prio), crate::sync::atomic::Ordering::AcqRel);
164    }
165
166    Some(crate::task::TaskId::new(prio, index))
167}
168
169/// Static cells so the waker data pointer has real provenance (miri
170/// strict-provenance clean) instead of an integer→pointer cast. 2 KiB of
171/// static storage; indexed by `(priority, index)`.
172static TASK_ID_CELLS: [[crate::task::TaskId; 32]; 32] = {
173    let mut cells = [[crate::task::TaskId::new(0, 0); 32]; 32];
174    let mut p = 0;
175    while p < 32 {
176        let mut i = 0;
177        while i < 32 {
178            cells[p][i] = crate::task::TaskId::new(p as u8, i as u8);
179            i += 1;
180        }
181        p += 1;
182    }
183    cells
184};
185
186fn encode_waker_data(id: crate::task::TaskId) -> *const () {
187    // SAFETY-free: `&TASK_ID_CELLS[...] as *const _` is a genuine pointer
188    // into a static; the cells are never mutated.
189    core::ptr::addr_of!(TASK_ID_CELLS[id.priority() as usize][id.index() as usize]) as *const ()
190}
191
192fn decode_waker_data(data: *const ()) -> crate::task::TaskId {
193    // SAFETY: the pointer was produced by `encode_waker_data` and still
194    // points into the immutable static cells.
195    unsafe { *(data as *const crate::task::TaskId) }
196}
197
198// ── RawWaker vtable ──────────────────────────────────────────────
199
200unsafe fn waker_clone(data: *const ()) -> RawWaker {
201    RawWaker::new(data, &WAKER_VTABLE)
202}
203
204unsafe fn waker_wake(data: *const ()) {
205    mark_ready(decode_waker_data(data));
206}
207
208unsafe fn waker_wake_by_ref(data: *const ()) {
209    mark_ready(decode_waker_data(data));
210}
211
212unsafe fn waker_drop(_data: *const ()) {
213    // Nothing to drop; data is statically allocated.
214}
215
216static WAKER_VTABLE: RawWakerVTable =
217    RawWakerVTable::new(waker_clone, waker_wake, waker_wake_by_ref, waker_drop);
218
219/// Kick every other hart so an executor idling in `wfi`/`waiti` on a
220/// different core notices new ready work. Exposed separately from
221/// [`wake_task`] so a caller waking several tasks in one batch (e.g.
222/// `timer::poll_timers` scanning every expired deadline inside one
223/// critical section) can broadcast once instead of once per task.
224///
225/// Safe to call from ISR/Handler-mode context — every board's periodic
226/// tick ISR already does, via `timer::poll_timers`.
227pub fn broadcast_reschedule() {
228    let hart = crate::port::arch::hart_id();
229    for other in 0..crate::config::MAX_HARTS {
230        if other != hart {
231            crate::port::arch::request_reschedule_on(other);
232        }
233    }
234}
235
236/// Mark `id` ready and broadcast a reschedule request to every other hart.
237///
238/// Use this — not the lower-level [`mark_ready`] — from any context
239/// *outside* the executor's own poll loop: an ISR, another hart, or a
240/// driver's completion handler. `mark_ready` alone only flips bitmap
241/// flags; on a single-hart build the interrupt that called this is enough
242/// to break the executor out of `wfi`, but on `RIVET_MAX_HARTS > 1` an
243/// executor idling on a *different* hart would never notice (this is the
244/// exact bug `timer::poll_timers`'s own broadcast fixed, found on real
245/// dual-core ESP32-S3 hardware via `smp_test.rs`). Safe to call from ISR
246/// context.
247pub fn wake_task(id: crate::task::TaskId) {
248    mark_ready(id);
249    broadcast_reschedule();
250}
251
252/// Create a `Waker` for the task identified by `id`.
253pub fn task_waker(id: crate::task::TaskId) -> Waker {
254    let data = encode_waker_data(id);
255    let raw = RawWaker::new(data, &WAKER_VTABLE);
256    // SAFETY: `raw` is built by [`task_waker`] from the static vtable and
257    // a pointer into the immutable `TASK_ID_CELLS` static with no drop
258    // state; `waker_drop` is a no-op, so transferring ownership into the
259    // `Waker` is safe.
260    unsafe { Waker::from_raw(raw) }
261}
262
263#[cfg(test)]
264mod tests {
265    use super::*;
266
267    #[test]
268    fn encode_decode_roundtrip() {
269        crate::kernel_test! {
270        let data = encode_waker_data(crate::task::TaskId::new(5, 13));
271        let id = decode_waker_data(data);
272        assert_eq!(id.priority(), 5);
273        assert_eq!(id.index(), 13);
274        }
275    }
276
277    #[test]
278    fn mark_and_dequeue_single() {
279        crate::kernel_test! {
280        mark_ready(crate::task::TaskId::new(2, 0));
281        assert_eq!(next_ready(), Some(crate::task::TaskId::new(2, 0)));
282        assert_eq!(next_ready(), None);
283        }
284    }
285
286    #[test]
287    fn priority_ordering() {
288        crate::kernel_test! {
289        mark_ready(crate::task::TaskId::new(1, 0));
290        mark_ready(crate::task::TaskId::new(5, 0));
291        mark_ready(crate::task::TaskId::new(3, 0));
292        assert_eq!(next_ready(), Some(crate::task::TaskId::new(5, 0)));
293        assert_eq!(next_ready(), Some(crate::task::TaskId::new(3, 0)));
294        assert_eq!(next_ready(), Some(crate::task::TaskId::new(1, 0)));
295        assert_eq!(next_ready(), None);
296        }
297    }
298
299    #[test]
300    fn multiple_tasks_same_priority() {
301        crate::kernel_test! {
302        mark_ready(crate::task::TaskId::new(3, 0));
303        mark_ready(crate::task::TaskId::new(3, 1));
304        mark_ready(crate::task::TaskId::new(3, 2));
305        assert_eq!(next_ready(), Some(crate::task::TaskId::new(3, 0)));
306        assert_eq!(next_ready(), Some(crate::task::TaskId::new(3, 1)));
307        assert_eq!(next_ready(), Some(crate::task::TaskId::new(3, 2)));
308        assert_eq!(next_ready(), None);
309        }
310    }
311
312    #[test]
313    fn interleaved_wake() {
314        crate::kernel_test! {
315        mark_ready(crate::task::TaskId::new(2, 0));
316        assert_eq!(next_ready(), Some(crate::task::TaskId::new(2, 0)));
317        assert_eq!(next_ready(), None);
318        mark_ready(crate::task::TaskId::new(2, 0));
319        mark_ready(crate::task::TaskId::new(4, 1));
320        assert_eq!(next_ready(), Some(crate::task::TaskId::new(4, 1)));
321        assert_eq!(next_ready(), Some(crate::task::TaskId::new(2, 0)));
322        assert_eq!(next_ready(), None);
323        }
324    }
325
326    #[test]
327    fn has_pending_detects_work() {
328        crate::kernel_test! {
329        assert!(!has_pending());
330        mark_ready(crate::task::TaskId::new(0, 0));
331        assert!(has_pending());
332        next_ready();
333        clear_pend();
334        assert!(!has_pending());
335        }
336    }
337}
rivet/waker.rs

rivet/
waker.rs
