Skip to main content

rusteal_core/
task.rs

1use std::panic::AssertUnwindSafe;
2use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
3use std::sync::{Mutex, OnceLock};
4use std::thread::ThreadId;
5use std::time::{Duration, Instant};
6
7use crate::lock_or_recover;
8
9type Job = Box<dyn FnOnce() + Send>;
10type TickFn = Box<dyn FnMut(f32) + Send>;
11
12static QUEUE: Mutex<Vec<Job>> = Mutex::new(Vec::new());
13static TICKERS: Mutex<Vec<(u64, TickFn)>> = Mutex::new(Vec::new());
14static REMOVED: Mutex<Vec<u64>> = Mutex::new(Vec::new());
15static NEXT_TICKER: AtomicU64 = AtomicU64::new(1);
16static RUNNING: AtomicUsize = AtomicUsize::new(0);
17static GAME_THREAD: OnceLock<ThreadId> = OnceLock::new();
18
19fn log(level: u8, message: &str) {
20    if crate::api::is_api_initialized() {
21        unsafe {
22            crate::ffi_dispatch::logging_log(level, message.as_ptr(), message.len() as u32);
23        }
24    }
25}
26
27pub fn is_game_thread() -> bool {
28    GAME_THREAD
29        .get()
30        .is_some_and(|id| *id == std::thread::current().id())
31}
32
33pub fn run_on_game_thread(f: impl FnOnce() + Send + 'static) {
34    lock_or_recover(&QUEUE).push(Box::new(f));
35}
36
37pub fn spawn<T: Send + 'static>(
38    work: impl FnOnce() -> T + Send + 'static,
39    done: impl FnOnce(T) + Send + 'static,
40) {
41    RUNNING.fetch_add(1, Ordering::SeqCst);
42
43    let started = std::thread::Builder::new()
44        .name("rusteal-task".into())
45        .spawn(move || {
46            let result = std::panic::catch_unwind(AssertUnwindSafe(work));
47
48            match result {
49                Ok(value) => run_on_game_thread(move || done(value)),
50                Err(_) => run_on_game_thread(|| {
51                    log(crate::LOG_ERROR, "[Rusteal] task::spawn: the work panicked")
52                }),
53            }
54
55            RUNNING.fetch_sub(1, Ordering::SeqCst);
56        });
57
58    if started.is_err() {
59        RUNNING.fetch_sub(1, Ordering::SeqCst);
60
61        log(
62            crate::LOG_ERROR,
63            "[Rusteal] task::spawn: could not start a thread",
64        );
65    }
66}
67
68#[derive(Debug, Clone, Copy, PartialEq, Eq)]
69#[must_use = "keep the handle to remove the hook, or drop it to keep the hook forever"]
70pub struct TickHandle(u64);
71
72impl TickHandle {
73    pub fn remove(self) {
74        let mut tickers = lock_or_recover(&TICKERS);
75        let before = tickers.len();
76        tickers.retain(|(id, _)| *id != self.0);
77
78        if tickers.len() == before {
79            lock_or_recover(&REMOVED).push(self.0);
80        }
81    }
82}
83
84pub fn on_tick(f: impl FnMut(f32) + Send + 'static) -> TickHandle {
85    let id = NEXT_TICKER.fetch_add(1, Ordering::Relaxed);
86    lock_or_recover(&TICKERS).push((id, Box::new(f)));
87    TickHandle(id)
88}
89
90#[doc(hidden)]
91pub fn mark_game_thread() {
92    let _ = GAME_THREAD.set(std::thread::current().id());
93}
94
95#[doc(hidden)]
96pub fn tick(delta_seconds: f32) {
97    mark_game_thread();
98
99    let jobs = std::mem::take(&mut *lock_or_recover(&QUEUE));
100
101    for job in jobs {
102        if std::panic::catch_unwind(AssertUnwindSafe(job)).is_err() {
103            log(
104                crate::LOG_ERROR,
105                "[Rusteal] a closure run on the game thread panicked",
106            );
107        }
108    }
109
110    let mut running = std::mem::take(&mut *lock_or_recover(&TICKERS));
111
112    for (_, hook) in running.iter_mut() {
113        if std::panic::catch_unwind(AssertUnwindSafe(|| hook(delta_seconds))).is_err() {
114            log(crate::LOG_ERROR, "[Rusteal] a tick hook panicked");
115        }
116    }
117
118    let removed = std::mem::take(&mut *lock_or_recover(&REMOVED));
119    running.retain(|(id, _)| !removed.contains(id));
120    let mut tickers = lock_or_recover(&TICKERS);
121    running.append(&mut tickers);
122    *tickers = running;
123}
124
125#[doc(hidden)]
126pub fn shutdown() {
127    let started = Instant::now();
128    let mut warned = false;
129
130    while RUNNING.load(Ordering::SeqCst) > 0 {
131        if !warned && started.elapsed() > Duration::from_secs(1) {
132            warned = true;
133
134            log(
135                crate::LOG_WARNING,
136                &format!(
137                    "[Rusteal] waiting for {} background task(s) before unloading",
138                    RUNNING.load(Ordering::SeqCst)
139                ),
140            );
141        }
142
143        std::thread::sleep(Duration::from_millis(1));
144    }
145
146    lock_or_recover(&QUEUE).clear();
147    lock_or_recover(&TICKERS).clear();
148    lock_or_recover(&REMOVED).clear();
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154    use std::sync::Arc;
155
156    #[test]
157    fn work_comes_back_on_the_ticking_thread() {
158        let seen = Arc::new(Mutex::new(Vec::<String>::new()));
159
160        let log = seen.clone();
161
162        spawn(
163            || 6 * 7,
164            move |answer| log.lock().unwrap().push(format!("done {answer}")),
165        );
166
167        let log = seen.clone();
168        let hook = on_tick(move |dt| log.lock().unwrap().push(format!("tick {dt}")));
169
170        let deadline = Instant::now() + Duration::from_secs(5);
171
172        while !seen.lock().unwrap().iter().any(|s| s == "done 42") {
173            assert!(Instant::now() < deadline, "the work never came back");
174            tick(0.5);
175            std::thread::sleep(Duration::from_millis(1));
176        }
177
178        assert!(is_game_thread());
179        assert!(seen.lock().unwrap().iter().any(|s| s == "tick 0.5"));
180
181        hook.remove();
182        seen.lock().unwrap().clear();
183        tick(0.25);
184        assert!(seen.lock().unwrap().is_empty(), "a removed hook ran");
185
186        let log = seen.clone();
187        run_on_game_thread(move || log.lock().unwrap().push("queued".into()));
188        shutdown();
189        tick(0.1);
190        assert!(seen.lock().unwrap().is_empty(), "shutdown kept the queue");
191    }
192}