Skip to main content

mlua_isle/
hook.rs

1//! Cancellation token and Lua debug hook.
2//!
3//! A [`CancelToken`] is a shared `AtomicBool` that can be checked from
4//! both Rust code and a Lua debug hook.  When cancelled, the debug hook
5//! raises a Lua error containing the sentinel `__isle_cancelled__`,
6//! which is recognized by [`IsleError::from(mlua::Error)`].
7
8use std::sync::atomic::{AtomicBool, Ordering};
9use std::sync::Arc;
10
11/// Cancellation signal shared between caller and Lua thread.
12///
13/// Clone is cheap (Arc).
14#[derive(Clone)]
15pub struct CancelToken {
16    flag: Arc<AtomicBool>,
17}
18
19impl CancelToken {
20    /// Create a new token (not cancelled).
21    pub fn new() -> Self {
22        Self {
23            flag: Arc::new(AtomicBool::new(false)),
24        }
25    }
26
27    /// Signal cancellation.
28    pub fn cancel(&self) {
29        self.flag.store(true, Ordering::Release);
30    }
31
32    /// Check whether cancellation has been requested.
33    pub fn is_cancelled(&self) -> bool {
34        self.flag.load(Ordering::Acquire)
35    }
36}
37
38impl Default for CancelToken {
39    fn default() -> Self {
40        Self::new()
41    }
42}
43
44/// Install a Lua debug hook that checks the cancel token every N instructions.
45///
46/// When the token is cancelled, the hook raises a Lua error with a
47/// sentinel message that [`IsleError`](crate::IsleError) recognizes as
48/// a cancellation.
49///
50/// # Instruction interval
51///
52/// The `interval` controls how often the check runs.  Lower values
53/// give faster cancellation response at the cost of overhead.
54/// A value of 1000 is a reasonable default.
55pub(crate) fn install_cancel_hook(
56    lua: &mlua::Lua,
57    token: CancelToken,
58    interval: u32,
59) -> Result<(), crate::IsleError> {
60    lua.set_hook(
61        mlua::HookTriggers::new().every_nth_instruction(interval),
62        move |_lua, _debug| {
63            if token.is_cancelled() {
64                Err(mlua::Error::runtime("__isle_cancelled__"))
65            } else {
66                Ok(mlua::VmState::Continue)
67            }
68        },
69    )
70    .map_err(crate::IsleError::from)
71}
72
73/// Remove the debug hook (restores normal execution speed).
74pub(crate) fn remove_hook(lua: &mlua::Lua) {
75    lua.remove_hook();
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81
82    #[test]
83    fn token_default_not_cancelled() {
84        let token = CancelToken::new();
85        assert!(!token.is_cancelled());
86    }
87
88    #[test]
89    fn token_cancel_sets_flag() {
90        let token = CancelToken::new();
91        let clone = token.clone();
92        token.cancel();
93        assert!(clone.is_cancelled());
94    }
95
96    #[test]
97    fn hook_interrupts_lua_loop() {
98        let lua = mlua::Lua::new();
99        let token = CancelToken::new();
100        install_cancel_hook(&lua, token.clone(), 100).unwrap();
101
102        // Schedule cancel after a short spin
103        let t = token.clone();
104        std::thread::spawn(move || {
105            std::thread::sleep(std::time::Duration::from_millis(10));
106            t.cancel();
107        });
108
109        let result: mlua::Result<()> = lua.load("while true do end").exec();
110        assert!(result.is_err());
111        let err_msg = result.unwrap_err().to_string();
112        assert!(
113            err_msg.contains("__isle_cancelled__"),
114            "expected cancellation sentinel, got: {err_msg}"
115        );
116    }
117}