diff --git a/mlua-sys/src/luau/lua.rs b/mlua-sys/src/luau/lua.rs index 66812d1f9..7ae838be0 100644 --- a/mlua-sys/src/luau/lua.rs +++ b/mlua-sys/src/luau/lua.rs @@ -591,6 +591,9 @@ pub struct lua_Debug { // These are shared between all coroutines. // +/// Type for interrupt callbacks invoked at Luau safepoints. +pub type lua_InterruptFunction = unsafe extern "C-unwind" fn(L: *mut lua_State, gc: c_int); + #[repr(C)] #[non_exhaustive] pub struct lua_Callbacks { @@ -598,7 +601,7 @@ pub struct lua_Callbacks { pub userdata: *mut c_void, /// gets called at safepoints (loop back edges, call/ret, gc) if set - pub interrupt: Option, + pub interrupt: Option, /// gets called when an unprotected error is raised (if longjmp is used) pub panic: Option, diff --git a/src/lib.rs b/src/lib.rs index f1b8acd86..8db3df753 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -146,7 +146,7 @@ pub use crate::debug::HookTriggers; #[cfg(any(feature = "luau", doc))] #[cfg_attr(docsrs, doc(cfg(feature = "luau")))] -pub use crate::{buffer::Buffer, vector::Vector}; +pub use crate::{buffer::Buffer, luau::InterruptHandle, vector::Vector}; #[cfg(feature = "serde")] #[doc(inline)] diff --git a/src/luau/interrupt.rs b/src/luau/interrupt.rs new file mode 100644 index 000000000..a3eb35dcc --- /dev/null +++ b/src/luau/interrupt.rs @@ -0,0 +1,95 @@ +use std::ptr; +use std::sync::Arc; + +use parking_lot::Mutex; + +use crate::Lua; + +/// A handle for arming an interrupt registered with [`Lua::set_interrupt_oneshot`]. +/// +/// Removing or replacing the interrupt, or destroying the VM, invalidates all clones. +/// Dropping a handle does not cancel pending requests. +#[derive(Clone, Debug)] +pub struct InterruptHandle { + inner: Arc>, +} + +#[derive(Debug)] +struct InterruptState { + #[cfg(feature = "luau")] + slot: *mut Option, + pending: bool, + running: bool, +} + +// SAFETY: Only the interrupt slot is accessed through this pointer. +unsafe impl Send for InterruptState {} + +impl InterruptHandle { + /// Requests one callback invocation at a Luau execution safepoint. + /// + /// Pending requests coalesce and are consumed before invocation, regardless of the callback's + /// outcome. Requests made during the callback are deferred until it exits. + /// + /// Returns `true` if accepted, including when already armed, or `false` if the handle is + /// invalidated. Later invalidation can cancel an accepted request. + pub fn arm(&self) -> bool { + let mut state = self.inner.lock(); + if state.slot.is_null() { + return false; + } + state.pending = true; + if !state.running { + unsafe { state.slot.write(Some(Lua::interrupt_proc)) }; + } + true + } + + // The caller must invalidate this registration before freeing the callback slot, and must + // invalidate any previous registration before constructing this one. + pub(crate) unsafe fn new(state: *mut ffi::lua_State) -> Self { + let slot = ptr::addr_of_mut!((*ffi::lua_callbacks(state)).interrupt); + slot.write(None); + Self { + inner: Arc::new(Mutex::new(InterruptState { + slot, + pending: false, + running: false, + })), + } + } + + pub(crate) fn invalidate(&self) { + let mut state = self.inner.lock(); + if !state.slot.is_null() { + unsafe { state.slot.write(None) }; + state.slot = ptr::null_mut(); + } + state.pending = false; + } + + pub(crate) fn begin(&self) -> Option { + let mut state = self.inner.lock(); + if state.slot.is_null() || !state.pending || state.running { + return None; + } + state.pending = false; + state.running = true; + unsafe { state.slot.write(None) }; + Some(InterruptGuard(self.clone())) + } +} + +#[must_use = "keep the guard alive until the interrupt callback finishes"] +pub(crate) struct InterruptGuard(InterruptHandle); + +impl Drop for InterruptGuard { + fn drop(&mut self) { + let mut state = self.0.inner.lock(); + state.running = false; + // A callback may have replaced or removed its own registration. + if state.pending && !state.slot.is_null() { + unsafe { state.slot.write(Some(Lua::interrupt_proc)) }; + } + } +} diff --git a/src/luau/mod.rs b/src/luau/mod.rs index b9813451a..c3b874365 100644 --- a/src/luau/mod.rs +++ b/src/luau/mod.rs @@ -17,6 +17,7 @@ use crate::traits::{FromLuaMulti, IntoLua}; use crate::types::MaybeSend; pub use heap_dump::HeapDump; +pub use interrupt::InterruptHandle; pub use require::{FsRequirer, NavigateError, Require}; #[cfg(feature = "luau-jit")] @@ -168,5 +169,6 @@ unsafe extern "C-unwind" fn lua_loadstring(state: *mut ffi::lua_State) -> c_int } mod heap_dump; +mod interrupt; mod json; mod require; diff --git a/src/prelude.rs b/src/prelude.rs index a80999604..38e6153c2 100644 --- a/src/prelude.rs +++ b/src/prelude.rs @@ -31,7 +31,7 @@ pub use crate::state::GcGenParams as LuaGcGenParams; #[cfg(any(feature = "luau", doc))] #[doc(no_inline)] pub use crate::{ - Buffer as LuaBuffer, Vector as LuaVector, + Buffer as LuaBuffer, InterruptHandle as LuaInterruptHandle, Vector as LuaVector, chunk::{CompileConstant as LuaCompileConstant, Compiler as LuaCompiler}, function::CoverageInfo as LuaCoverageInfo, luau::{ diff --git a/src/state.rs b/src/state.rs index 6b2e7ea9d..eee3907e5 100644 --- a/src/state.rs +++ b/src/state.rs @@ -36,7 +36,7 @@ use crate::value::{Nil, Value}; use crate::{debug::HookTriggers, types::HookKind}; #[cfg(any(feature = "luau", doc))] -use crate::{buffer::Buffer, chunk::Compiler}; +use crate::{buffer::Buffer, chunk::Compiler, luau::InterruptHandle}; #[cfg(feature = "async")] use { @@ -796,6 +796,7 @@ impl Lua { /// Any Luau code is guaranteed to call this handler "eventually" /// (in practice this can happen at any function call or at any loop iteration). /// This is similar to `Lua::set_hook` but in more simplified form. + /// Replaces any previous interrupt registration, invalidating its one-shot handles. /// /// The provided interrupt function can error, and this error will be propagated through /// the Luau code that was executing at the time the interrupt was triggered. @@ -844,39 +845,63 @@ impl Lua { where F: Fn(&Lua) -> Result + MaybeSend + 'static, { - unsafe extern "C-unwind" fn interrupt_proc(state: *mut ffi::lua_State, gc: c_int) { - if gc >= 0 { - // We don't support GC interrupts since they cannot survive Lua exceptions - return; - } - let result = callback_error_ext(state, ptr::null_mut(), false, move |extra, _| { - let interrupt_cb = (*extra).interrupt_callback.clone(); - let interrupt_cb = mlua_expect!(interrupt_cb, "no interrupt callback set in interrupt_proc"); - if XRc::strong_count(&interrupt_cb) > 2 { - return Ok(VmState::Continue); // Don't allow recursion - } - interrupt_cb((*extra).lua()) - }); - match result { - VmState::Continue => {} - VmState::Yield => { - // We can yield only at yieldable points, otherwise ignore and continue - if ffi::lua_isyieldable(state) != 0 { - ffi::lua_yield(state, 0); - } - } + // Set interrupt callback + let lua = self.lock(); + unsafe { + if let Some(handle) = (*lua.extra.get()).interrupt_handle.take() { + handle.invalidate(); } + let _old_callback = (*lua.extra.get()).interrupt_callback.replace(XRc::new(callback)); + (*ffi::lua_callbacks(lua.main_state())).interrupt = Some(Self::interrupt_proc); } + } - // Set interrupt callback + /// Registers a reusable interrupt callback triggered by [`InterruptHandle::arm`]. + /// + /// The interrupt is initially disarmed and automatically disarms before each invocation. + /// While disarmed, no native interrupt callback is installed. + /// + /// The callback runs on the thread executing Luau, with the error and yield behavior of + /// [`Lua::set_interrupt`]. Replaces any previous interrupt registration, invalidating its + /// handles. + /// + /// # Example + /// + /// ``` + /// # use mlua::{Error, Lua}; + /// # #[cfg(feature = "luau")] + /// # fn main() -> mlua::Result<()> { + /// let lua = Lua::new(); + /// let interrupt = lua.set_interrupt_oneshot(|_| Err(Error::runtime("interrupted"))); + /// // A timer or watchdog can arm the interrupt from another thread. + /// std::thread::spawn(move || interrupt.arm()).join().unwrap(); + /// assert!(lua.load("while true do end").exec().is_err()); + /// # Ok(()) + /// # } + /// # #[cfg(not(feature = "luau"))] + /// # fn main() {} + /// ``` + #[cfg(any(feature = "luau", doc))] + #[cfg_attr(docsrs, doc(cfg(feature = "luau")))] + #[must_use] + pub fn set_interrupt_oneshot(&self, callback: F) -> InterruptHandle + where + F: Fn(&Lua) -> Result + MaybeSend + 'static, + { let lua = self.lock(); unsafe { + if let Some(handle) = (*lua.extra.get()).interrupt_handle.take() { + handle.invalidate(); + } let _old_callback = (*lua.extra.get()).interrupt_callback.replace(XRc::new(callback)); - (*ffi::lua_callbacks(lua.main_state())).interrupt = Some(interrupt_proc); + let handle = InterruptHandle::new(lua.main_state()); + (*lua.extra.get()).interrupt_handle = Some(handle.clone()); + handle } } - /// Removes any interrupt function previously set by `set_interrupt`. + /// Removes any interrupt function previously set by [`Lua::set_interrupt`] or + /// [`Lua::set_interrupt_oneshot`], invalidating its handles. /// /// This function has no effect if an 'interrupt' was not previously set. #[cfg(any(feature = "luau", doc))] @@ -884,11 +909,47 @@ impl Lua { pub fn remove_interrupt(&self) { let lua = self.lock(); unsafe { + if let Some(handle) = (*lua.extra.get()).interrupt_handle.take() { + handle.invalidate(); + } let _old_callback = (*lua.extra.get()).interrupt_callback.take(); (*ffi::lua_callbacks(lua.main_state())).interrupt = None; } } + #[cfg(feature = "luau")] + pub(crate) unsafe extern "C-unwind" fn interrupt_proc(state: *mut ffi::lua_State, gc: c_int) { + if gc >= 0 { + // We don't support GC interrupts since they cannot survive Lua exceptions. + return; + } + let result = callback_error_ext(state, ptr::null_mut(), false, |extra, _| { + // Reserve error storage before consuming a one-shot request. + let _guard = match (*extra).interrupt_handle.as_ref() { + Some(handle) => match handle.begin() { + Some(guard) => Some(guard), + None => return Ok(VmState::Continue), + }, + None => None, + }; + let interrupt_cb = (*extra).interrupt_callback.clone(); + let interrupt_cb = mlua_expect!(interrupt_cb, "no interrupt callback set in interrupt_proc"); + if XRc::strong_count(&interrupt_cb) > 2 { + return Ok(VmState::Continue); // Don't allow recursion. + } + interrupt_cb((*extra).lua()) + }); + match result { + VmState::Continue => {} + VmState::Yield => { + // We can yield only at yieldable points, otherwise ignore and continue. + if ffi::lua_isyieldable(state) != 0 { + ffi::lua_yield(state, 0); + } + } + } + } + /// Sets a callback invoked when thread lifecycle events occur. /// /// `triggers` controls which events trigger the callback, see [`ThreadTriggers`] for more diff --git a/src/state/extra.rs b/src/state/extra.rs index 038855d3a..c2fef27d5 100644 --- a/src/state/extra.rs +++ b/src/state/extra.rs @@ -90,6 +90,8 @@ pub(crate) struct ExtraData { pub(super) warn_callback: Option, #[cfg(feature = "luau")] pub(super) interrupt_callback: Option, + #[cfg(feature = "luau")] + pub(super) interrupt_handle: Option, pub(super) thread_triggers: ThreadTriggers, pub(super) thread_event_callback: Option, pub(super) thread_event_state: *mut ffi::lua_State, @@ -108,6 +110,10 @@ pub(crate) struct ExtraData { impl Drop for ExtraData { fn drop(&mut self) { + #[cfg(feature = "luau")] + if let Some(handle) = self.interrupt_handle.take() { + handle.invalidate(); + } unsafe { if !self.owned { self.lua.assume_init_drop(); @@ -206,6 +212,8 @@ impl ExtraData { warn_callback: None, #[cfg(feature = "luau")] interrupt_callback: None, + #[cfg(feature = "luau")] + interrupt_handle: None, thread_triggers: ThreadTriggers::default(), thread_event_callback: None, thread_event_state: ptr::null_mut(), diff --git a/src/state/raw.rs b/src/state/raw.rs index 8060414a6..2e9c3c6db 100644 --- a/src/state/raw.rs +++ b/src/state/raw.rs @@ -71,6 +71,9 @@ impl Drop for RawLua { #[cfg(feature = "luau")] { + if let Some(handle) = (*self.extra.get()).interrupt_handle.take() { + handle.invalidate(); + } // Reset any callbacks (*ffi::lua_callbacks(self.main_state())).interrupt = None; (*ffi::lua_callbacks(self.main_state())).userthread = None; diff --git a/tests/luau.rs b/tests/luau.rs index 0ae10cc8b..cbed00c26 100644 --- a/tests/luau.rs +++ b/tests/luau.rs @@ -5,7 +5,10 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use mlua::chunk::Compiler; -use mlua::{Error, Function, Lua, LuaOptions, ObjectLike, Result, StdLib, Table, Value, Vector, VmState}; +use mlua::{ + Error, Function, InterruptHandle, Lua, LuaOptions, ObjectLike, Result, StdLib, Table, Value, Vector, + VmState, ffi, +}; #[test] fn test_version() -> Result<()> { @@ -365,20 +368,263 @@ fn test_interrupt_callback_drop() { } } - for replace in [false, true] { + for oneshot in [false, true] { + for replacement in 0..3 { + let lua = Lua::new(); + let guard = ReenterOnDrop(lua.clone()); + let callback = move |_: &Lua| { + let _ = &guard; + panic!("called a callback during drop"); + }; + let old = if oneshot { + let handle = lua.set_interrupt_oneshot(callback); + handle.arm(); + Some(handle) + } else { + lua.set_interrupt(callback); + None + }; + match replacement { + 0 => lua.remove_interrupt(), + 1 => lua.set_interrupt(|_| Ok(VmState::Continue)), + _ => { + let _handle = lua.set_interrupt_oneshot(|_| panic!("unarmed replacement")); + } + } + assert_eq!(lua.remove_app_data::<()>(), Some(())); + if let Some(old) = old { + assert!(!old.arm()); + } + } + } +} + +static_assertions::assert_impl_all!(InterruptHandle: Clone, Send, Sync); + +fn has_native_interrupt(lua: &Lua) -> Result { + let mut present = false; + unsafe { + lua.exec_raw::<()>((), |state| { + present = (*ffi::lua_callbacks(state)).interrupt.is_some(); + })?; + } + Ok(present) +} + +#[test] +fn test_interrupt_oneshot() -> Result<()> { + let lua = Lua::new(); + + let f = lua.load("for i = 1, 10 do end").into_function()?; + let calls = Arc::new(AtomicU64::new(0)); + let handle = lua.set_interrupt_oneshot({ + let calls = calls.clone(); + move |_| { + calls.fetch_add(1, Ordering::Relaxed); + Ok(VmState::Continue) + } + }); + assert!(!has_native_interrupt(&lua)?); + f.call::<()>(())?; + assert_eq!(calls.load(Ordering::Relaxed), 0); + + assert!(handle.arm()); + assert!(handle.clone().arm()); + // GC safepoints must not consume the pending execution interrupt. + unsafe { + lua.exec_raw::<()>((), |state| { + (*ffi::lua_callbacks(state)).interrupt.unwrap()(state, 0); + })?; + } + assert_eq!(calls.load(Ordering::Relaxed), 0); + f.call::<()>(())?; + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert!(!has_native_interrupt(&lua)?); + f.call::<()>(())?; + assert_eq!(calls.load(Ordering::Relaxed), 1); + + // Rearming is reusable, and dropping the handle does not cancel a pending request. + assert!(handle.arm()); + drop(handle); + f.call::<()>(())?; + assert_eq!(calls.load(Ordering::Relaxed), 2); + + Ok(()) +} + +#[test] +fn test_interrupt_oneshot_rearm() -> Result<()> { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + #[derive(Clone, Copy)] + enum Outcome { + Continue, + Error, + Panic, + } + + // Exercise guard cleanup on normal return, error, and panic. + for outcome in [Outcome::Continue, Outcome::Error, Outcome::Panic] { let lua = Lua::new(); - let guard = ReenterOnDrop(lua.clone()); - lua.set_interrupt(move |_| { - let _ = &guard; - panic!("called a callback during drop"); + let f = lua.load("for i = 1, 10 do end").into_function()?; + let calls = Arc::new(AtomicU64::new(0)); + let handle = lua.set_interrupt_oneshot({ + let calls = calls.clone(); + move |lua| { + if calls.fetch_add(1, Ordering::Relaxed) == 0 { + assert!(lua.app_data_ref::().unwrap().arm()); + assert!(!has_native_interrupt(lua)?); + lua.load("for i = 1, 10 do end").exec()?; + assert_eq!(calls.load(Ordering::Relaxed), 1); // No recursive delivery. + match outcome { + Outcome::Continue => {} + Outcome::Error => return Err(Error::runtime("interrupt error")), + Outcome::Panic => panic!("interrupt panic"), + } + } + Ok(VmState::Continue) + } }); - if replace { - lua.set_interrupt(|_| Ok(VmState::Continue)); - } else { - lua.remove_interrupt(); + lua.set_app_data(handle.clone()); + handle.arm(); + let result = catch_unwind(AssertUnwindSafe(|| f.call::<()>(()))); + match outcome { + Outcome::Continue => result.unwrap()?, + Outcome::Error => { + assert!(matches!(result.unwrap(), Err(Error::RuntimeError(msg)) if msg == "interrupt error")) + } + Outcome::Panic => assert!(result.is_err()), } - assert_eq!(lua.remove_app_data::<()>(), Some(())); + f.call::<()>(())?; + assert_eq!(calls.load(Ordering::Relaxed), 2); + assert!(!has_native_interrupt(&lua)?); } + + Ok(()) +} + +#[test] +fn test_interrupt_oneshot_replace_during_callback() -> Result<()> { + for persistent in [false, true] { + let lua = Lua::new(); + let handle = lua.set_interrupt_oneshot(move |lua| { + assert!(lua.app_data_ref::().unwrap().arm()); + if persistent { + lua.set_interrupt(|_| Ok(VmState::Continue)); + } else { + let _handle = lua.set_interrupt_oneshot(|_| panic!("unarmed replacement")); + } + Ok(VmState::Continue) + }); + lua.set_app_data(handle.clone()); + handle.arm(); + lua.load("for i = 1, 10 do end").exec()?; + assert!(!handle.arm()); + // Finishing the old callback must not change the replacement's native slot. + assert_eq!(has_native_interrupt(&lua)?, persistent); + } + Ok(()) +} + +#[test] +fn test_interrupt_oneshot_yield() -> Result<()> { + let lua = Lua::new(); + + let co = lua.create_thread(lua.load("for i = 1, 10 do end; return 42").into_function()?)?; + let handle = lua.set_interrupt_oneshot(|_| Ok(VmState::Yield)); + handle.arm(); + co.resume::<()>(())?; + assert!(co.is_resumable()); + assert!(!has_native_interrupt(&lua)?); + assert_eq!(co.resume::(())?, 42); + assert!(co.is_finished()); + + // A request is also consumed when yielding is not possible. + let f = lua.create_function(|lua, ()| lua.load("return 42").eval::())?; + handle.arm(); + assert_eq!(f.call::(())?, 42); + assert!(!has_native_interrupt(&lua)?); + + Ok(()) +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn test_interrupt_oneshot_from_another_thread() -> Result<()> { + use std::sync::mpsc; + use std::time::Duration; + + let lua = Lua::new(); + + let (tx, rx) = mpsc::channel(); + lua.globals().set( + "started", + lua.create_function(move |_, ()| { + tx.send(()).unwrap(); + Ok(()) + })?, + )?; + let f = lua + .load("started(); local n = 0; for i = 1, 100000000 do n += i end; return n") + .into_function()?; + let handle = lua.set_interrupt_oneshot(|_| Err(Error::runtime("interrupted"))); + let worker = std::thread::spawn(move || { + rx.recv_timeout(Duration::from_secs(10)).unwrap(); + // Let the VM enter its bounded loop, so missed interrupts fail instead of hanging. + std::thread::sleep(Duration::from_millis(1)); + assert!(handle.arm()); + }); + let result = f.call::(()); + worker.join().unwrap(); + assert!(matches!(result, Err(Error::RuntimeError(msg)) if msg == "interrupted")); + + Ok(()) +} + +#[cfg(not(target_family = "wasm"))] +#[test] +fn test_interrupt_oneshot_teardown_race() { + use std::sync::Barrier; + + for remove in [false, true] { + for _ in 0..64 { + let lua = Lua::new(); + let handle = lua.set_interrupt_oneshot(|_| Ok(VmState::Continue)); + let barrier = Arc::new(Barrier::new(2)); + let worker = std::thread::spawn({ + let barrier = barrier.clone(); + move || { + barrier.wait(); + while handle.arm() { + std::thread::yield_now(); + } + } + }); + barrier.wait(); + if remove { + lua.remove_interrupt(); + } + drop(lua); + worker.join().unwrap(); + } + } +} + +#[cfg(not(feature = "send"))] +#[test] +fn test_interrupt_oneshot_non_send_callback() { + use std::rc::Rc; + + let lua = Lua::new(); + let captured = Rc::new(()); + let weak = Rc::downgrade(&captured); + let handle = lua.set_interrupt_oneshot(move |_| { + let _ = &captured; + Ok(VmState::Continue) + }); + drop(lua); + assert!(weak.upgrade().is_none()); // Handles must not retain the callback or VM. + assert!(!handle.arm()); } #[test]