vj_sim/
fiber.rs

1//! Stackful fibers for blocking DPI import tasks (LRM 1800-2023 35.5.3,
2//! 35.7, 35.9).
3//!
4//! A DPI `import task` is a C function that may call an exported SV task
5//! which consumes time. The C frame cannot be unwound into a stackless
6//! process state machine, so the call runs on a [`Fiber`]: a private stack
7//! the scheduler can leave ([`yield_now`], from inside the fiber) and later
8//! re-enter ([`Fiber::resume`]). A call that never blocks finishes on its
9//! first `resume` and costs one stack switch pair.
10//!
11//! Built on glibc's `makecontext`/`swapcontext` (declared here directly: no
12//! `libc` dependency in this workspace). `ucontext_t` is kept as an opaque,
13//! over-sized buffer; the only fields touched go through those functions,
14//! except `uc_stack` / `uc_link`, which are located by the x86_64 / aarch64
15//! glibc layout checked in the tests.
16//!
17//! A panic never crosses the fiber boundary: the entry closure runs under
18//! `catch_unwind` and the payload is re-raised by [`Fiber::resume`] on the
19//! resuming stack.
20
21use std::any::Any;
22use std::cell::Cell;
23use std::panic::{AssertUnwindSafe, catch_unwind, resume_unwind};
24
25/// Bytes reserved for a `ucontext_t` (936 on x86_64 glibc, 4560 on aarch64).
26const UCONTEXT_BYTES: usize = 5120;
27/// Fiber stack size. A whole process kernel can run on a fiber (a design with a
28/// timed export task, `CompiledKernelProcess::with_task_fibers`), so it is as
29/// deep as the kernels need; the pages are committed only when touched.
30const STACK_BYTES: usize = 4 << 20;
31/// Stacks kept for reuse after a fiber is dropped (a mapping costs more than a
32/// short call), per thread.
33const POOL_MAX: usize = 32;
34const GUARD_BYTES: usize = 4096;
35
36#[repr(C, align(64))]
37struct UContext([u8; UCONTEXT_BYTES]);
38
39unsafe extern "C" {
40    fn getcontext(ucp: *mut UContext) -> i32;
41    fn makecontext(ucp: *mut UContext, func: extern "C" fn(), argc: i32, ...);
42    fn swapcontext(oucp: *mut UContext, ucp: *const UContext) -> i32;
43    fn mmap(addr: *mut u8, len: usize, prot: i32, flags: i32, fd: i32, off: i64) -> *mut u8;
44    fn munmap(addr: *mut u8, len: usize) -> i32;
45    fn mprotect(addr: *mut u8, len: usize, prot: i32) -> i32;
46}
47
48const PROT_NONE: i32 = 0;
49const PROT_RW: i32 = 3;
50const MAP_PRIVATE_ANON_STACK: i32 = 0x02 | 0x20 | 0x20000 | 0x4000; // PRIVATE | ANONYMOUS | STACK | NORESERVE
51const MAP_FAILED: *mut u8 = !0usize as *mut u8;
52
53/// Whether a [`Fiber::resume`] left the fiber parked or ran it to its end.
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
55pub enum FiberState {
56    /// The fiber called [`yield_now`]; resume it later.
57    Yielded,
58    /// The entry closure returned; the fiber is spent.
59    Finished,
60}
61
62struct Inner {
63    /// The fiber's saved context while it is parked.
64    fiber_ctx: UContext,
65    /// The resumer's saved context while the fiber runs.
66    caller_ctx: UContext,
67    entry: Option<Box<dyn FnOnce()>>,
68    panic: Option<Box<dyn Any + Send>>,
69    done: bool,
70    map: *mut u8,
71}
72
73/// One suspendable C call.
74pub struct Fiber {
75    inner: Box<Inner>,
76}
77
78thread_local! {
79    /// Unused stack mappings (guard page included), most recent last.
80    static POOL: std::cell::RefCell<Vec<usize>> = const { std::cell::RefCell::new(Vec::new()) };
81    /// The fiber currently running on this thread (null on the main stack).
82    static CURRENT: Cell<*mut Inner> = const { Cell::new(std::ptr::null_mut()) };
83}
84
85extern "C" fn trampoline() {
86    let inner = CURRENT.with(Cell::get);
87    // SAFETY: `resume` set CURRENT to the live `Inner` before switching here.
88    let f = unsafe { (*inner).entry.take() };
89    if let Some(f) = f
90        && let Err(p) = catch_unwind(AssertUnwindSafe(f))
91    {
92        // SAFETY: as above.
93        unsafe { (*inner).panic = Some(p) };
94    }
95    // SAFETY: as above. Returning to the resumer; never resumed again.
96    unsafe {
97        (*inner).done = true;
98        swapcontext(&raw mut (*inner).fiber_ctx, &raw const (*inner).caller_ctx);
99    }
100    unreachable!("a finished fiber was resumed");
101}
102
103impl Fiber {
104    /// A fiber that will run `entry` on its own stack on the first
105    /// [`Self::resume`]. `None` when the stack cannot be mapped.
106    pub fn new(entry: impl FnOnce() + 'static) -> Option<Fiber> {
107        let total = STACK_BYTES + GUARD_BYTES;
108        let pooled = POOL.with(|p| p.borrow_mut().pop());
109        let map = match pooled {
110            Some(m) => m as *mut u8,
111            None => {
112                // SAFETY: plain anonymous mapping; checked for failure.
113                let map = unsafe {
114                    mmap(
115                        std::ptr::null_mut(),
116                        total,
117                        PROT_RW,
118                        MAP_PRIVATE_ANON_STACK,
119                        -1,
120                        0,
121                    )
122                };
123                if map == MAP_FAILED {
124                    return None;
125                }
126                // SAFETY: the lowest page of the mapping, made a guard.
127                unsafe { mprotect(map, GUARD_BYTES, PROT_NONE) };
128                map
129            }
130        };
131        let mut inner = Box::new(Inner {
132            fiber_ctx: UContext([0; UCONTEXT_BYTES]),
133            caller_ctx: UContext([0; UCONTEXT_BYTES]),
134            entry: Some(Box::new(entry)),
135            panic: None,
136            done: false,
137            map,
138        });
139        // SAFETY: `fiber_ctx` is a zeroed, aligned buffer larger than
140        // `ucontext_t`; `set_stack` writes the documented stack fields.
141        unsafe {
142            if getcontext(&raw mut inner.fiber_ctx) != 0 {
143                release_stack(map);
144                return None;
145            }
146            set_stack(
147                &mut inner.fiber_ctx,
148                map.add(GUARD_BYTES),
149                STACK_BYTES,
150                std::ptr::null_mut(),
151            );
152            makecontext(&raw mut inner.fiber_ctx, trampoline, 0);
153        }
154        Some(Fiber { inner })
155    }
156
157    /// Runs the fiber until it yields or finishes. Resumes a panic raised
158    /// inside it. Must not be called from inside the same fiber.
159    pub fn resume(&mut self) -> FiberState {
160        assert!(!self.inner.done, "resume of a finished fiber");
161        let inner: *mut Inner = &raw mut *self.inner;
162        let prev = CURRENT.with(|c| c.replace(inner));
163        // SAFETY: both contexts live in the boxed `Inner`.
164        unsafe { swapcontext(&raw mut (*inner).caller_ctx, &raw const (*inner).fiber_ctx) };
165        CURRENT.with(|c| c.set(prev));
166        if let Some(p) = self.inner.panic.take() {
167            resume_unwind(p);
168        }
169        if self.inner.done {
170            FiberState::Finished
171        } else {
172            FiberState::Yielded
173        }
174    }
175
176    /// Whether the entry closure has returned.
177    pub fn is_finished(&self) -> bool {
178        self.inner.done
179    }
180}
181
182impl Drop for Fiber {
183    fn drop(&mut self) {
184        // A parked fiber's stack is simply discarded: its frames never run
185        // their destructors (an abandoned blocked import task, e.g. at the
186        // end of simulation). Nothing on it is owned by anyone else.
187        release_stack(self.inner.map);
188    }
189}
190
191/// Returns a stack mapping made in [`Fiber::new`], unused from now on, to the
192/// pool (or to the OS when the pool is full).
193fn release_stack(map: *mut u8) {
194    let kept = POOL.with(|p| {
195        let mut p = p.borrow_mut();
196        let room = p.len() < POOL_MAX;
197        if room {
198            p.push(map as usize);
199        }
200        room
201    });
202    if !kept {
203        // SAFETY: the mapping made in `new`, unused once the fiber is gone.
204        unsafe { munmap(map, STACK_BYTES + GUARD_BYTES) };
205    }
206}
207
208/// Parks the running fiber and returns to its resumer. `false` (nothing
209/// happened) when called on the main stack.
210pub fn yield_now() -> bool {
211    let inner = CURRENT.with(Cell::get);
212    if inner.is_null() {
213        return false;
214    }
215    // SAFETY: CURRENT is the live fiber that is calling us.
216    unsafe { swapcontext(&raw mut (*inner).fiber_ctx, &raw const (*inner).caller_ctx) };
217    true
218}
219
220/// Whether the caller runs on a fiber (a blocking export is legal here).
221pub fn on_fiber() -> bool {
222    !CURRENT.with(Cell::get).is_null()
223}
224
225/// Writes `uc_link`, `uc_stack.{ss_sp,ss_size}` of a glibc `ucontext_t`.
226///
227/// glibc x86_64 and aarch64 share the head layout
228/// `{ unsigned long uc_flags; ucontext_t *uc_link; stack_t uc_stack; ... }`
229/// with `stack_t = { void *ss_sp; int ss_flags; size_t ss_size; }`.
230///
231/// # Safety
232/// `uc` must have been filled by `getcontext`.
233unsafe fn set_stack(uc: &mut UContext, sp: *mut u8, size: usize, link: *mut UContext) {
234    let w = uc.0.as_mut_ptr().cast::<usize>();
235    // SAFETY: offsets 1..=4 words lie within the buffer.
236    unsafe {
237        w.add(1).write(link as usize); // uc_link
238        w.add(2).write(sp as usize); // ss_sp
239        w.add(3).write(0); // ss_flags (+ padding)
240        w.add(4).write(size); // ss_size
241    }
242}
243
244#[cfg(test)]
245mod tests {
246    use super::*;
247    use std::rc::Rc;
248
249    #[test]
250    fn runs_to_completion_without_yielding() {
251        let hit = Rc::new(Cell::new(0));
252        let h = hit.clone();
253        let mut f = Fiber::new(move || h.set(7)).unwrap();
254        assert!(!on_fiber());
255        assert_eq!(f.resume(), FiberState::Finished);
256        assert_eq!(hit.get(), 7);
257        assert!(f.is_finished());
258    }
259
260    #[test]
261    fn yields_and_resumes_in_order() {
262        let log = Rc::new(std::cell::RefCell::new(Vec::new()));
263        let l = log.clone();
264        let mut f = Fiber::new(move || {
265            assert!(on_fiber());
266            l.borrow_mut().push(1);
267            assert!(yield_now());
268            l.borrow_mut().push(3);
269            yield_now();
270            l.borrow_mut().push(5);
271        })
272        .unwrap();
273        assert_eq!(f.resume(), FiberState::Yielded);
274        log.borrow_mut().push(2);
275        assert_eq!(f.resume(), FiberState::Yielded);
276        log.borrow_mut().push(4);
277        assert_eq!(f.resume(), FiberState::Finished);
278        assert_eq!(*log.borrow(), [1, 2, 3, 4, 5]);
279    }
280
281    #[test]
282    fn interleaved_fibers_keep_their_own_stacks() {
283        let out = Rc::new(std::cell::RefCell::new(Vec::new()));
284        let mk = |tag: i32| {
285            let o = out.clone();
286            Fiber::new(move || {
287                let local = [tag; 64];
288                yield_now();
289                o.borrow_mut().push(local.iter().sum::<i32>());
290            })
291            .unwrap()
292        };
293        let (mut a, mut b) = (mk(1), mk(2));
294        a.resume();
295        b.resume();
296        // Out-of-order completion: the case a nested scheduler cannot do.
297        assert_eq!(b.resume(), FiberState::Finished);
298        assert_eq!(a.resume(), FiberState::Finished);
299        assert_eq!(*out.borrow(), [128, 64]);
300    }
301
302    #[test]
303    fn yield_on_main_stack_is_a_no_op() {
304        assert!(!yield_now());
305    }
306
307    #[test]
308    fn panic_is_re_raised_on_the_resumer() {
309        let mut f = Fiber::new(|| panic!("boom")).unwrap();
310        let r = catch_unwind(AssertUnwindSafe(|| f.resume()));
311        assert!(r.is_err());
312    }
313
314    #[test]
315    fn nested_fibers() {
316        let log = Rc::new(std::cell::RefCell::new(Vec::new()));
317        let l = log.clone();
318        let mut outer = Fiber::new(move || {
319            let l2 = l.clone();
320            let mut inner = Fiber::new(move || {
321                l2.borrow_mut().push("in1");
322                yield_now();
323                l2.borrow_mut().push("in2");
324            })
325            .unwrap();
326            inner.resume();
327            l.borrow_mut().push("out1");
328            yield_now();
329            inner.resume();
330            l.borrow_mut().push("out2");
331        })
332        .unwrap();
333        outer.resume();
334        outer.resume();
335        assert_eq!(*log.borrow(), ["in1", "out1", "in2", "out2"]);
336    }
337}
vj_sim/fiber.rs

vj_sim/
fiber.rs
