Add one-shot Luau interrupts

mlua-sys/src/luau/lua.rs

@@ -591,14 +591,17 @@ 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 {
     /// arbitrary userdata pointer that is never overwritten by Luau
     pub userdata: *mut c_void,
 
     /// gets called at safepoints (loop back edges, call/ret, gc) if set
-    pub interrupt: Option<unsafe extern "C-unwind" fn(L: *mut lua_State, gc: c_int)>,
+    pub interrupt: Option<lua_InterruptFunction>,
     /// gets called when an unprotected error is raised (if longjmp is used)
     pub panic: Option<unsafe extern "C-unwind" fn(L: *mut lua_State, errcode: c_int)>,
 

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)]

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<Mutex<InterruptState>>,
+}
+
+#[derive(Debug)]
+struct InterruptState {
+    #[cfg(feature = "luau")]
+    slot: *mut Option<ffi::lua_InterruptFunction>,
+    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<InterruptGuard> {
+        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)) };
+        }
+    }
+}

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;

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::{

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,51 +845,111 @@ impl Lua {
     where
         F: Fn(&Lua) -> Result<VmState> + 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<F>(&self, callback: F) -> InterruptHandle
+    where
+        F: Fn(&Lua) -> Result<VmState> + 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))]
     #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
     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

src/state/extra.rs

@@ -90,6 +90,8 @@ pub(crate) struct ExtraData {
     pub(super) warn_callback: Option<crate::types::WarnCallback>,
     #[cfg(feature = "luau")]
     pub(super) interrupt_callback: Option<crate::types::InterruptCallback>,
+    #[cfg(feature = "luau")]
+    pub(super) interrupt_handle: Option<crate::luau::InterruptHandle>,
     pub(super) thread_triggers: ThreadTriggers,
     pub(super) thread_event_callback: Option<ThreadEventCallback>,
     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(),

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;