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/// Install a cancel hook on a Lua coroutine thread.
79///
80/// Same as [`install_cancel_hook`] but targets a specific [`Thread`](mlua::Thread)
81/// instead of the main Lua state.  Used for cooperative coroutine execution
82/// where each coroutine needs its own cancel check.
83#[cfg(feature = "tokio")]
84pub(crate) fn install_cancel_hook_on_thread(
85    thread: &mlua::Thread,
86    token: CancelToken,
87    interval: u32,
88) -> Result<(), crate::IsleError> {
89    thread
90        .set_hook(
91            mlua::HookTriggers::new().every_nth_instruction(interval),
92            move |_lua, _debug| {
93                if token.is_cancelled() {
94                    Err(mlua::Error::runtime("__isle_cancelled__"))
95                } else {
96                    Ok(mlua::VmState::Continue)
97                }
98            },
99        )
100        .map_err(crate::IsleError::from)
101}
102
103#[cfg(test)]
104mod tests {
105    use super::*;
106
107    #[test]
108    fn token_default_not_cancelled() {
109        let token = CancelToken::new();
110        assert!(!token.is_cancelled());
111    }
112
113    #[test]
114    fn token_cancel_sets_flag() {
115        let token = CancelToken::new();
116        let clone = token.clone();
117        token.cancel();
118        assert!(clone.is_cancelled());
119    }
120
121    #[test]
122    fn hook_interrupts_lua_loop() {
123        let lua = mlua::Lua::new();
124        let token = CancelToken::new();
125        install_cancel_hook(&lua, token.clone(), 100).unwrap();
126
127        // Schedule cancel after a short spin
128        let t = token.clone();
129        std::thread::spawn(move || {
130            std::thread::sleep(std::time::Duration::from_millis(10));
131            t.cancel();
132        });
133
134        let result: mlua::Result<()> = lua.load("while true do end").exec();
135        assert!(result.is_err());
136        let err_msg = result.unwrap_err().to_string();
137        assert!(
138            err_msg.contains("__isle_cancelled__"),
139            "expected cancellation sentinel, got: {err_msg}"
140        );
141    }
142}
mlua_isle/hook.rs

mlua_isle/
hook.rs
