rusteal_core/
task.rs

1// Work off the game thread, and the way back to it.
2//
3// UE's objects belong to the game thread: a background thread must not call
4// the engine. `spawn` runs work on a thread of its own and hands its result
5// to the game thread, `run_on_game_thread` queues a closure from anywhere, and
6// `on_tick` runs one every frame. The plugin's core ticker drives them (the
7// library's `on_tick` callback), in the editor and in games alike.
8
9use std::panic::AssertUnwindSafe;
10use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
11use std::sync::{Mutex, OnceLock};
12use std::thread::ThreadId;
13use std::time::{Duration, Instant};
14
15use crate::lock_or_recover;
16
17type Job = Box<dyn FnOnce() + Send>;
18type TickFn = Box<dyn FnMut(f32) + Send>;
19
20static QUEUE: Mutex<Vec<Job>> = Mutex::new(Vec::new());
21static TICKERS: Mutex<Vec<(u64, TickFn)>> = Mutex::new(Vec::new());
22/// Tick hooks removed while the hooks run; dropped when they are put back.
23static REMOVED: Mutex<Vec<u64>> = Mutex::new(Vec::new());
24static NEXT_TICKER: AtomicU64 = AtomicU64::new(1);
25/// Threads `spawn` started that have not finished their work.
26static RUNNING: AtomicUsize = AtomicUsize::new(0);
27static GAME_THREAD: OnceLock<ThreadId> = OnceLock::new();
28
29fn log(level: u8, message: &str) {
30    if crate::api::is_api_initialized() {
31        // SAFETY: the API table is initialized.
32        unsafe {
33            crate::ffi_dispatch::logging_log(level, message.as_ptr(), message.len() as u32);
34        }
35    }
36}
37
38/// Whether the caller is on the game thread, where UE may be called.
39pub fn is_game_thread() -> bool {
40    GAME_THREAD.get().is_some_and(|id| *id == std::thread::current().id())
41}
42
43/// Run `f` on the game thread at the next frame: from a background thread,
44/// the way back to the engine.
45pub fn run_on_game_thread(f: impl FnOnce() + Send + 'static) {
46    lock_or_recover(&QUEUE).push(Box::new(f));
47}
48
49/// Run `work` on a thread of its own, then `done` with its result on the game
50/// thread. `work` must not call the engine; `done` may.
51///
52/// The library waits for every `work` still running before it unloads (a hot
53/// reload, the engine closing), so no thread runs code that is gone; a `done`
54/// still queued then is dropped.
55pub fn spawn<T: Send + 'static>(
56    work: impl FnOnce() -> T + Send + 'static,
57    done: impl FnOnce(T) + Send + 'static,
58) {
59    RUNNING.fetch_add(1, Ordering::SeqCst);
60    let started = std::thread::Builder::new().name("rusteal-task".into()).spawn(move || {
61        let result = std::panic::catch_unwind(AssertUnwindSafe(work));
62        match result {
63            Ok(value) => run_on_game_thread(move || done(value)),
64            Err(_) => run_on_game_thread(|| {
65                log(crate::LOG_ERROR, "[Rusteal] task::spawn: the work panicked")
66            }),
67        }
68        RUNNING.fetch_sub(1, Ordering::SeqCst);
69    });
70    if started.is_err() {
71        RUNNING.fetch_sub(1, Ordering::SeqCst);
72        log(crate::LOG_ERROR, "[Rusteal] task::spawn: could not start a thread");
73    }
74}
75
76/// A hook `on_tick` added; `remove` stops it. Dropping the handle keeps the
77/// hook running.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79#[must_use = "keep the handle to remove the hook, or drop it to keep the hook forever"]
80pub struct TickHandle(u64);
81
82impl TickHandle {
83    /// Stop calling the hook.
84    pub fn remove(self) {
85        let mut tickers = lock_or_recover(&TICKERS);
86        let before = tickers.len();
87        tickers.retain(|(id, _)| *id != self.0);
88        if tickers.len() == before {
89            // Running right now: dropped when the hooks are put back.
90            lock_or_recover(&REMOVED).push(self.0);
91        }
92    }
93}
94
95/// Call `f` every frame on the game thread with the frame's seconds, in the
96/// editor too, with or without a world.
97pub fn on_tick(f: impl FnMut(f32) + Send + 'static) -> TickHandle {
98    let id = NEXT_TICKER.fetch_add(1, Ordering::Relaxed);
99    lock_or_recover(&TICKERS).push((id, Box::new(f)));
100    TickHandle(id)
101}
102
103/// Record the calling thread as the game thread. Called by the runtime's
104/// init, which runs on it.
105#[doc(hidden)]
106pub fn mark_game_thread() {
107    let _ = GAME_THREAD.set(std::thread::current().id());
108}
109
110/// The frame: the queued closures, then the tick hooks. Called by the
111/// runtime's `on_tick` callback.
112#[doc(hidden)]
113pub fn tick(delta_seconds: f32) {
114    mark_game_thread();
115
116    let jobs = std::mem::take(&mut *lock_or_recover(&QUEUE));
117    for job in jobs {
118        if std::panic::catch_unwind(AssertUnwindSafe(job)).is_err() {
119            log(crate::LOG_ERROR, "[Rusteal] a closure run on the game thread panicked");
120        }
121    }
122
123    // Out of the lock while they run: a hook may add or remove hooks.
124    let mut running = std::mem::take(&mut *lock_or_recover(&TICKERS));
125    for (_, hook) in running.iter_mut() {
126        if std::panic::catch_unwind(AssertUnwindSafe(|| hook(delta_seconds))).is_err() {
127            log(crate::LOG_ERROR, "[Rusteal] a tick hook panicked");
128        }
129    }
130    let removed = std::mem::take(&mut *lock_or_recover(&REMOVED));
131    running.retain(|(id, _)| !removed.contains(id));
132    let mut tickers = lock_or_recover(&TICKERS);
133    running.append(&mut tickers);
134    *tickers = running;
135}
136
137/// Before the library unloads: wait for the spawned work, then forget the
138/// queue and the hooks. Called by the runtime's `on_shutdown` callback.
139#[doc(hidden)]
140pub fn shutdown() {
141    let started = Instant::now();
142    let mut warned = false;
143    while RUNNING.load(Ordering::SeqCst) > 0 {
144        if !warned && started.elapsed() > Duration::from_secs(1) {
145            warned = true;
146            log(crate::LOG_WARNING, &format!(
147                "[Rusteal] waiting for {} background task(s) before unloading",
148                RUNNING.load(Ordering::SeqCst)
149            ));
150        }
151        std::thread::sleep(Duration::from_millis(1));
152    }
153    lock_or_recover(&QUEUE).clear();
154    lock_or_recover(&TICKERS).clear();
155    lock_or_recover(&REMOVED).clear();
156}
157
158#[cfg(test)]
159mod tests {
160    use super::*;
161    use std::sync::Arc;
162
163    // One test: the queues are process-wide.
164    #[test]
165    fn work_comes_back_on_the_ticking_thread() {
166        let seen = Arc::new(Mutex::new(Vec::<String>::new()));
167
168        let log = seen.clone();
169        spawn(|| 6 * 7, move |answer| log.lock().unwrap().push(format!("done {answer}")));
170        let log = seen.clone();
171        let hook = on_tick(move |dt| log.lock().unwrap().push(format!("tick {dt}")));
172
173        let deadline = Instant::now() + Duration::from_secs(5);
174        while !seen.lock().unwrap().iter().any(|s| s == "done 42") {
175            assert!(Instant::now() < deadline, "the work never came back");
176            tick(0.5);
177            std::thread::sleep(Duration::from_millis(1));
178        }
179        assert!(is_game_thread());
180        assert!(seen.lock().unwrap().iter().any(|s| s == "tick 0.5"));
181
182        hook.remove();
183        seen.lock().unwrap().clear();
184        tick(0.25);
185        assert!(seen.lock().unwrap().is_empty(), "a removed hook ran");
186
187        let log = seen.clone();
188        run_on_game_thread(move || log.lock().unwrap().push("queued".into()));
189        shutdown();
190        tick(0.1);
191        assert!(seen.lock().unwrap().is_empty(), "shutdown kept the queue");
192    }
193}
rusteal_core/task.rs

rusteal_core/
task.rs
