gasm/
thread.rs

1//! Cooperative threads for games with their own loop (design/threads.md), the Rust
2//! side of the C SDK's `gasm_thread.h`: the same scheduler, the same rules.
3//!
4//! Every thread runs on the guest's one wasm thread, until it blocks, yields or
5//! waits for the next frame; then the next ready thread runs, in creation order. A
6//! frame ends when no thread is ready. So there are no data races, and the schedule
7//! depends only on the input: runs stay reproducible on every runner.
8//!
9//! Threads need [`threaded_main_loop!`](crate::threaded_main_loop!) (and Binaryen's
10//! `wasm-opt --asyncify`, as `main_loop!`); [`spawn`] panics elsewhere. The locks
11//! in [`crate::sync`] work everywhere.
12//!
13//! ```ignore
14//! fn run() -> i32 {
15//!     let loader = gasm::thread::spawn(|| load_level("e1m1"));   // may wait for frames
16//!     while !loader.is_finished() { draw_loading_screen(); gasm::main_loop::wait_frame(); }
17//!     let level = loader.join();
18//!     ...
19//! }
20//! gasm::threaded_main_loop!(run);
21//! ```
22//!
23//! How it works: each thread is suspended with Asyncify into its own buffer and has
24//! its own shadow stack. Asyncify leaves the wasm stack pointer alone both ways (an
25//! unwind skips the functions' epilogues, a rewind their prologues), so a suspended
26//! thread's frames are everything above the stack pointer it had when it unwound:
27//! the scheduler records it then, sets it back before rewinding the thread, and runs
28//! its own calls on a stack of its own (`gasm__get_sp`/`gasm__set_sp`: sp/gasm_sp.s).
29//!
30//! Differences from `std::thread`: closures needn't be `Send` (it is one wasm
31//! thread), `thread_local!` values are shared by all threads, time is the frame's
32//! ([`sleep`] ends at the first frame at or after its deadline), and a thread that
33//! spins without blocking or yielding keeps the frame from ending.
34
35#![cfg_attr(not(target_arch = "wasm32"), allow(dead_code, unused_imports))]
36
37use std::cell::{Cell, RefCell, UnsafeCell};
38use std::rc::Rc;
39use std::time::Duration;
40
41/// A thread's default stack (its locals, the shadow stack) and Asyncify buffer.
42pub const DEFAULT_STACK: usize = 256 * 1024;
43const ASYNC_SIZE: usize = 256 * 1024;
44const MAIN_ASYNC_SIZE: usize = 1024 * 1024;
45const SCHED_STACK: usize = 64 * 1024;
46/// "No deadline" for [`block`].
47pub const FOREVER: f64 = -1.0;
48
49#[cfg(target_arch = "wasm32")]
50mod asm {
51    #[link(wasm_import_module = "asyncify")]
52    unsafe extern "C" {
53        pub fn start_unwind(data: *mut u8);
54        pub fn stop_unwind();
55        pub fn start_rewind(data: *mut u8);
56        pub fn stop_rewind();
57    }
58    // sp/gasm_sp.s (libgasm_sp.a, linked by build.rs)
59    unsafe extern "C" {
60        pub fn gasm__get_sp() -> usize;
61        pub fn gasm__set_sp(sp: usize);
62    }
63}
64
65#[derive(Clone, Copy, PartialEq, Eq, Debug)]
66enum State {
67    Ready,
68    Blocked,
69    WaitFrame,
70    Done,
71}
72
73#[repr(C)]
74struct AsyncData {
75    cur: *mut u8,
76    end: *mut u8,
77}
78
79struct Thread {
80    id: u64,
81    state: State,
82    body: Option<Box<dyn FnOnce()>>,
83    started: bool,
84    detached: bool,
85    timed_out: bool,
86    /// its stack pointer where it unwound (suspended)
87    sp: usize,
88    /// Blocked: what it waits for (0: only a deadline)
89    on: usize,
90    /// Blocked: the time_ms() when it times out, or FOREVER
91    deadline: f64,
92    /// its stack (empty for main: the module's own)
93    stack: Vec<u128>,
94    adata: AsyncData,
95    abuf: Vec<u8>,
96}
97
98impl Thread {
99    fn new(id: u64, body: Option<Box<dyn FnOnce()>>, stack: usize, abuf: usize) -> Box<Thread> {
100        Box::new(Thread {
101            id,
102            state: State::Ready,
103            body,
104            started: false,
105            detached: false,
106            timed_out: false,
107            sp: 0,
108            on: 0,
109            deadline: FOREVER,
110            stack: vec![0u128; stack.div_ceil(16)],
111            adata: AsyncData { cur: std::ptr::null_mut(), end: std::ptr::null_mut() },
112            abuf: vec![0u8; abuf],
113        })
114    }
115    fn top(&self) -> usize {
116        self.stack.as_ptr_range().end as usize
117    }
118}
119
120struct Sched {
121    /// creation order; [0] is main. Boxed: Asyncify keeps a pointer to a suspended
122    /// thread's `adata` while spawns grow the Vec
123    #[allow(clippy::vec_box)]
124    threads: Vec<Box<Thread>>,
125    current: usize,
126    cursor: usize,
127    next_id: u64,
128    active: bool,
129    started: bool,
130    finished: bool,
131    main: Option<fn() -> i32>,
132    main_result: i32,
133    main_base: usize,
134    frames: u32,
135}
136
137// Asyncify's state, read right after an unwind returns and in switch_impl during a
138// rewind: plain statics, no calls (nothing instrumented may run in between).
139static mut REWINDING: bool = false;
140static mut UNWINDING: bool = false;
141
142struct Global(UnsafeCell<Option<Sched>>);
143// One wasm thread; only the frame export and the threads it runs touch it.
144unsafe impl Sync for Global {}
145static SCHED: Global = Global(UnsafeCell::new(None));
146
147#[repr(align(16))]
148struct SchedStack(UnsafeCell<[u8; SCHED_STACK]>);
149unsafe impl Sync for SchedStack {}
150static SCHED_STACK_BUF: SchedStack = SchedStack(UnsafeCell::new([0; SCHED_STACK]));
151
152fn sched() -> &'static mut Sched {
153    // single-threaded: no two of these references are used at the same time
154    unsafe {
155        (*SCHED.0.get()).get_or_insert_with(|| Sched {
156            threads: Vec::new(),
157            current: 0,
158            cursor: 0,
159            next_id: 1,
160            active: false,
161            started: false,
162            finished: false,
163            main: None,
164            main_result: 0,
165            main_base: 0,
166            frames: 0,
167        })
168    }
169}
170
171fn sched_top() -> usize {
172    SCHED_STACK_BUF.0.get() as usize + SCHED_STACK
173}
174
175/// Suspend the current thread (its state already set) to the scheduler; returns
176/// when the scheduler resumes it. Reached through a pointer (Asyncify instruments
177/// the callers of indirect calls, while this stays plain, like an import).
178fn switch_impl() {
179    #[cfg(target_arch = "wasm32")]
180    unsafe {
181        if REWINDING {
182            asm::stop_rewind();
183            REWINDING = false;
184            return;
185        }
186        let s = sched();
187        let t = &mut s.threads[s.current];
188        t.adata.cur = t.abuf.as_mut_ptr();
189        t.adata.end = t.adata.cur.add(t.abuf.len());
190        UNWINDING = true;
191        asm::start_unwind((&raw mut t.adata).cast::<u8>());
192    }
193}
194static SWITCH: Global2 = Global2(switch_impl);
195struct Global2(fn());
196
197fn suspend(state: State) {
198    let s = sched();
199    let cur = s.current;
200    s.threads[cur].state = state;
201    let f = unsafe { std::ptr::read_volatile(&SWITCH.0) };
202    f()
203}
204
205/// Whether cooperative threads run (inside `threaded_main_loop!`).
206pub fn available() -> bool {
207    sched().active
208}
209
210/// Wait until [`wake`] is called on `on` or `deadline_ms` ([`crate::time_ms`] time,
211/// [`FOREVER`]: none) passes; true if it timed out. For building locks: `on` is any
212/// address that identifies what is waited for. Without the scheduler, a wait with a
213/// deadline times out at once and one without is a deadlock (trap).
214pub fn block(on: usize, deadline_ms: f64) -> bool {
215    let s = sched();
216    if !s.active {
217        if deadline_ms >= 0.0 {
218            return true;
219        }
220        crate::log("gasm::thread: deadlock (the only thread waits for something no other thread can do)");
221        std::process::abort();
222    }
223    let cur = s.current;
224    let t = &mut s.threads[cur];
225    t.on = on;
226    t.deadline = deadline_ms;
227    t.timed_out = false;
228    suspend(State::Blocked);
229    sched().threads[sched().current].timed_out
230}
231
232/// Make threads blocked on `on` ready (the first one in creation order, or all).
233pub fn wake(on: usize, one: bool) {
234    if on == 0 {
235        return;
236    }
237    for t in sched().threads.iter_mut() {
238        if t.state == State::Blocked && t.on == on {
239            t.state = State::Ready;
240            if one {
241                return;
242            }
243        }
244    }
245}
246
247/// Let the other ready threads run first.
248pub fn yield_now() {
249    if sched().active {
250        suspend(State::Ready);
251    }
252}
253
254/// Wait at least `d` of frame time (other threads run meanwhile).
255pub fn sleep(d: Duration) {
256    let ms = d.as_secs_f64() * 1000.0;
257    if ms <= 0.0 {
258        yield_now();
259    } else {
260        block(0, crate::time_ms() + ms);
261    }
262}
263
264/// Wait for the next frame (the same as `main_loop::wait_frame` in a threaded loop).
265pub fn wait_frame() {
266    suspend(State::WaitFrame);
267}
268
269/// A thread's number: main is 1, then in creation order.
270#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, PartialOrd, Ord)]
271pub struct ThreadId(pub u64);
272
273/// The running thread's id.
274pub fn current() -> ThreadId {
275    let s = sched();
276    ThreadId(s.threads.get(s.current).map_or(1, |t| t.id))
277}
278
279/// The handle of a spawned thread: [`join`](JoinHandle::join) waits for its result.
280/// Dropping it detaches the thread (it runs to the end on its own).
281pub struct JoinHandle<T> {
282    id: u64,
283    result: Rc<RefCell<Option<T>>>,
284    joined: Cell<bool>,
285}
286
287impl<T> JoinHandle<T> {
288    pub fn thread_id(&self) -> ThreadId {
289        ThreadId(self.id)
290    }
291    /// The thread has returned (join won't wait).
292    pub fn is_finished(&self) -> bool {
293        self.result.borrow().is_some()
294    }
295    /// Wait for the thread to finish and take its result.
296    pub fn join(self) -> T {
297        while self.result.borrow().is_none() {
298            block(thread_addr(self.id), FOREVER);
299        }
300        self.joined.set(true);
301        let r = self.result.borrow_mut().take().expect("result");
302        // the scheduler frees a finished thread once nobody can join it
303        if let Some(t) = sched().threads.iter_mut().find(|t| t.id == self.id) {
304            t.detached = true;
305        }
306        free_done();
307        r
308    }
309}
310
311impl<T> Drop for JoinHandle<T> {
312    fn drop(&mut self) {
313        if !self.joined.get() {
314            if let Some(t) = sched().threads.iter_mut().find(|t| t.id == self.id) {
315                t.detached = true;
316            }
317            free_done();
318        }
319    }
320}
321
322/// What joiners block on: a thread's id (never 0).
323fn thread_addr(id: u64) -> usize {
324    usize::MAX - id as usize
325}
326
327/// Remove finished, detached threads (never the running one or main).
328fn free_done() {
329    let s = sched();
330    let cur_id = s.threads.get(s.current).map(|t| t.id);
331    let cursor_id = s.threads.get(s.cursor).map(|t| t.id);
332    s.threads.retain(|t| t.id == 1 || !(t.state == State::Done && t.detached) || Some(t.id) == cur_id);
333    s.current = cur_id.and_then(|id| s.threads.iter().position(|t| t.id == id)).unwrap_or(0);
334    s.cursor = cursor_id.and_then(|id| s.threads.iter().position(|t| t.id == id)).unwrap_or(0);
335}
336
337/// Thread options (the stack size).
338#[derive(Default)]
339pub struct Builder {
340    stack: Option<usize>,
341}
342
343impl Builder {
344    pub fn new() -> Builder {
345        Builder::default()
346    }
347    /// The thread's stack (its locals; default [`DEFAULT_STACK`]).
348    pub fn stack_size(mut self, bytes: usize) -> Builder {
349        self.stack = Some(bytes);
350        self
351    }
352    /// Start a thread running `f`; it first runs when the current thread blocks,
353    /// yields or waits for a frame.
354    pub fn spawn<F, T>(self, f: F) -> JoinHandle<T>
355    where
356        F: FnOnce() -> T + 'static,
357        T: 'static,
358    {
359        let s = sched();
360        if !s.active {
361            panic!("gasm::thread::spawn: threads need gasm::threaded_main_loop! (and wasm-opt --asyncify)");
362        }
363        let result = Rc::new(RefCell::new(None));
364        let out = result.clone();
365        let id = s.next_id;
366        s.next_id += 1;
367        let body: Box<dyn FnOnce()> = Box::new(move || {
368            let v = f();
369            *out.borrow_mut() = Some(v);
370        });
371        s.threads.push(Thread::new(id, Some(body), self.stack.unwrap_or(DEFAULT_STACK), ASYNC_SIZE));
372        JoinHandle { id, result, joined: Cell::new(false) }
373    }
374}
375
376/// Start a thread running `f` (see [`Builder::spawn`]).
377pub fn spawn<F, T>(f: F) -> JoinHandle<T>
378where
379    F: FnOnce() -> T + 'static,
380    T: 'static,
381{
382    Builder::new().spawn(f)
383}
384
385/// Frames the threaded loop has run.
386pub fn frames() -> u32 {
387    sched().frames
388}
389
390#[doc(hidden)]
391pub fn __set_main(f: fn() -> i32) {
392    sched().main = Some(f);
393}
394
395// A thread's body; instrumented (an unwind returns from here) and not inlined.
396// By id: other threads may come and go while it is suspended.
397#[inline(never)]
398fn run_thread(id: u64) {
399    if id == 1 {
400        let main = sched().main;
401        let rc = main.map_or(0, |f| f());
402        if unsafe { UNWINDING } {
403            return;
404        }
405        sched().main_result = rc;
406    } else {
407        let body = sched().threads.iter_mut().find(|t| t.id == id).and_then(|t| t.body.take());
408        if let Some(b) = body {
409            b();
410        }
411        if unsafe { UNWINDING } {
412            return;
413        }
414    }
415    if let Some(t) = sched().threads.iter_mut().find(|t| t.id == id) {
416        t.state = State::Done;
417    }
418}
419
420/// The next ready thread after the cursor, in creation order.
421fn pick(s: &Sched) -> Option<usize> {
422    let n = s.threads.len();
423    (1..=n).map(|k| (s.cursor + k) % n).find(|&i| s.threads[i].state == State::Ready)
424}
425
426/// One frame: run ready threads until all of them wait (for the next frame, a
427/// deadline or each other). Inlined into the `gasm_loop_frame` export, which is on
428/// Asyncify's remove list: it never unwinds.
429#[doc(hidden)]
430#[inline(always)]
431pub fn __threaded_loop_frame() {
432    #[cfg(target_arch = "wasm32")]
433    unsafe {
434        let s = sched();
435        if s.finished {
436            return;
437        }
438        let entry_sp = asm::gasm__get_sp();
439        if !s.started {
440            s.started = true;
441            s.active = true;
442            s.main_base = entry_sp;
443            s.threads.push(Thread::new(s.next_id, None, 0, MAIN_ASYNC_SIZE));
444            s.next_id += 1;
445        } else {
446            s.frames += 1;
447        }
448        asm::gasm__set_sp(sched_top());
449        let now = crate::time_ms();
450        for t in sched().threads.iter_mut() {
451            if t.state == State::WaitFrame {
452                t.state = State::Ready;
453            } else if t.state == State::Blocked && t.deadline >= 0.0 && now >= t.deadline {
454                t.state = State::Ready;
455                t.timed_out = true;
456            }
457        }
458        while let Some(i) = pick(sched()) {
459            let s = sched();
460            s.cursor = i;
461            s.current = i;
462            let t = &mut s.threads[i];
463            // a new thread starts at its base; a suspended one continues where it unwound
464            let sp = if t.started { t.sp } else if i == 0 { s.main_base } else { t.top() };
465            let resume = t.started;
466            t.started = true;
467            let id = t.id;
468            asm::gasm__set_sp(sp);
469            if resume {
470                REWINDING = true;
471                asm::start_rewind((&raw mut s.threads[i].adata).cast::<u8>());
472            }
473            run_thread(id);
474            if UNWINDING {
475                asm::stop_unwind();
476                UNWINDING = false;
477                let unwound_sp = asm::gasm__get_sp();
478                let s = sched();
479                s.threads[s.current].sp = unwound_sp;
480            }
481            asm::gasm__set_sp(sched_top());
482            let s = sched();
483            let t = &s.threads[s.current];
484            if t.state == State::Done {
485                if t.id == 1 {
486                    s.finished = true;
487                    let code = s.main_result;
488                    asm::gasm__set_sp(entry_sp);
489                    crate::exit(code); // proc_exit: the runner ends the game
490                }
491                let id = t.id;
492                wake(thread_addr(id), false); // joiners
493                free_done();
494            }
495        }
496        let s = sched();
497        s.current = 0;
498        // the frame ends; if nothing can ever run again, say who waits for what
499        let alive = s.threads.iter().any(|t| t.state == State::WaitFrame || (t.state == State::Blocked && t.deadline >= 0.0));
500        if !alive {
501            crate::log("gasm::thread: deadlock: every thread is blocked");
502            for t in &s.threads {
503                if t.state == State::Blocked {
504                    crate::log(&format!("  thread {} waits on {:#x}", t.id, t.on));
505                }
506            }
507            std::process::abort();
508        }
509        asm::gasm__set_sp(entry_sp);
510    }
511    #[cfg(not(target_arch = "wasm32"))]
512    {
513        // natively (the stub host) there is only main, run to the end
514        let s = sched();
515        if s.finished {
516            return;
517        }
518        s.finished = true;
519        let rc = s.main.map_or(0, |f| f());
520        crate::exit(rc);
521    }
522}
gasm/thread.rs

gasm/
thread.rs
