Support user thread creation callback (Luau)

khvzak · khvzak · commit 1f970824b4f8 · 2025-03-28T11:51:35.000Z
diff --git a/src/state.rs b/src/state.rs
@@ -419,7 +419,7 @@ impl Lua {
         // Make sure that Lua is initialized
         let _ = Self::get_or_init_from_ptr(state);
 
-        callback_error_ext(state, ptr::null_mut(), move |extra, nargs| {
+        callback_error_ext(state, ptr::null_mut(), true, move |extra, nargs| {
             let rawlua = (*extra).raw_lua();
             let _guard = StateGuard::new(rawlua, state);
             let args = A::from_stack_args(nargs, 1, None, rawlua)?;
@@ -652,7 +652,7 @@ impl Lua {
                 // We don't support GC interrupts since they cannot survive Lua exceptions
                 return;
             }
-            let result = callback_error_ext(state, ptr::null_mut(), move |extra, _| {
+            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 {
@@ -690,6 +690,60 @@ impl Lua {
         }
     }
 
+    /// Sets a thread event callback that will be called when a thread is created or destroyed.
+    ///
+    /// The callback is called with a [`Value`] argument that is either:
+    /// - A [`Thread`] object when thread is created
+    /// - A [`LightUserData`] when thread is destroyed
+    #[cfg(any(feature = "luau", doc))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
+    pub fn set_thread_event_callback<F>(&self, callback: F)
+    where
+        F: Fn(&Lua, Value) -> Result<()> + MaybeSend + 'static,
+    {
+        unsafe extern "C-unwind" fn userthread_proc(parent: *mut ffi::lua_State, child: *mut ffi::lua_State) {
+            let extra = ExtraData::get(child);
+            let thread_cb = match (*extra).userthread_callback {
+                Some(ref cb) => cb.clone(),
+                None => return,
+            };
+            if XRc::strong_count(&thread_cb) > 2 {
+                return; // Don't allow recursion
+            }
+            let value = match parent.is_null() {
+                // Thread is about to be destroyed, pass light userdata
+                true => Value::LightUserData(crate::LightUserData(child as _)),
+                false => {
+                    // Thread is created, pass thread object
+                    ffi::lua_pushthread(child);
+                    ffi::lua_xmove(child, (*extra).ref_thread, 1);
+                    Value::Thread(Thread((*extra).raw_lua().pop_ref_thread(), child))
+                }
+            };
+            callback_error_ext((*extra).raw_lua().state(), extra, false, move |extra, _| {
+                thread_cb((*extra).lua(), value)
+            })
+        }
+
+        // Set thread callback
+        let lua = self.lock();
+        unsafe {
+            (*lua.extra.get()).userthread_callback = Some(XRc::new(callback));
+            (*ffi::lua_callbacks(lua.main_state())).userthread = Some(userthread_proc);
+        }
+    }
+
+    /// Removes any thread event callback previously set by `set_thread_event_callback`.
+    #[cfg(any(feature = "luau", doc))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
+    pub fn remove_thread_event_callback(&self) {
+        let lua = self.lock();
+        unsafe {
+            (*lua.extra.get()).userthread_callback = None;
+            (*ffi::lua_callbacks(lua.main_state())).userthread = None;
+        }
+    }
+
     /// Sets the warning function to be used by Lua to emit warnings.
     ///
     /// Requires `feature = "lua54"`
@@ -705,7 +759,7 @@ impl Lua {
 
         unsafe extern "C-unwind" fn warn_proc(ud: *mut c_void, msg: *const c_char, tocont: c_int) {
             let extra = ud as *mut ExtraData;
-            callback_error_ext((*extra).raw_lua().state(), extra, |extra, _| {
+            callback_error_ext((*extra).raw_lua().state(), extra, false, |extra, _| {
                 let warn_callback = (*extra).warn_callback.clone();
                 let warn_callback = mlua_expect!(warn_callback, "no warning callback set in warn_proc");
                 if XRc::strong_count(&warn_callback) > 2 {
@@ -1444,7 +1498,7 @@ impl Lua {
                     Err(_) => return,
                 },
                 ffi::LUA_TTHREAD => {
-                    ffi::lua_newthread(state);
+                    ffi::lua_pushthread(state);
                 }
                 #[cfg(feature = "luau")]
                 ffi::LUA_TBUFFER => {
diff --git a/src/state/extra.rs b/src/state/extra.rs
@@ -80,6 +80,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) userthread_callback: Option<crate::types::UserThreadCallback>,
 
     #[cfg(feature = "luau")]
     pub(super) sandboxed: bool,
@@ -177,6 +179,8 @@ impl ExtraData {
             #[cfg(feature = "luau")]
             interrupt_callback: None,
             #[cfg(feature = "luau")]
+            userthread_callback: None,
+            #[cfg(feature = "luau")]
             sandboxed: false,
             #[cfg(feature = "luau")]
             compiler: None,
diff --git a/src/state/raw.rs b/src/state/raw.rs
@@ -70,6 +70,13 @@ impl Drop for RawLua {
 
             let mem_state = MemoryState::get(self.main_state());
 
+            #[cfg(feature = "luau")]
+            {
+                // Reset any callbacks
+                (*ffi::lua_callbacks(self.main_state())).interrupt = None;
+                (*ffi::lua_callbacks(self.main_state())).userthread = None;
+            }
+
             ffi::lua_close(self.main_state());
 
             // Deallocate `MemoryState`
@@ -420,7 +427,7 @@ impl RawLua {
         }
 
         unsafe extern "C-unwind" fn global_hook_proc(state: *mut ffi::lua_State, ar: *mut ffi::lua_Debug) {
-            let status = callback_error_ext(state, ptr::null_mut(), move |extra, _| {
+            let status = callback_error_ext(state, ptr::null_mut(), false, move |extra, _| {
                 match (*extra).hook_callback.clone() {
                     Some(hook_callback) => {
                         let rawlua = (*extra).raw_lua();
@@ -453,7 +460,7 @@ impl RawLua {
                 return;
             }
 
-            let status = callback_error_ext(state, ptr::null_mut(), |extra, _| {
+            let status = callback_error_ext(state, ptr::null_mut(), false, |extra, _| {
                 let rawlua = (*extra).raw_lua();
                 let _guard = StateGuard::new(rawlua, state);
                 let debug = Debug::new(rawlua, ar);
@@ -564,7 +571,11 @@ impl RawLua {
         let _sg = StackGuard::new(state);
         check_stack(state, 3)?;
 
-        let thread_state = if self.unlikely_memory_error() {
+        let protect = !self.unlikely_memory_error();
+        #[cfg(feature = "luau")]
+        let protect = protect || (*self.extra.get()).userthread_callback.is_some();
+
+        let thread_state = if !protect {
             ffi::lua_newthread(state)
         } else {
             protect_lua!(state, 0, 1, |state| ffi::lua_newthread(state))?
@@ -1177,7 +1188,7 @@ impl RawLua {
     pub(crate) fn create_callback(&self, func: Callback) -> Result<Function> {
         unsafe extern "C-unwind" fn call_callback(state: *mut ffi::lua_State) -> c_int {
             let upvalue = get_userdata::<CallbackUpvalue>(state, ffi::lua_upvalueindex(1));
-            callback_error_ext(state, (*upvalue).extra.get(), |extra, nargs| {
+            callback_error_ext(state, (*upvalue).extra.get(), true, |extra, nargs| {
                 // Lua ensures that `LUA_MINSTACK` stack spaces are available (after pushing arguments)
                 // The lock must be already held as the callback is executed
                 let rawlua = (*extra).raw_lua();
@@ -1226,7 +1237,7 @@ impl RawLua {
             // so the first upvalue is always valid
             let upvalue = get_userdata::<AsyncCallbackUpvalue>(state, ffi::lua_upvalueindex(1));
             let extra = (*upvalue).extra.get();
-            callback_error_ext(state, extra, |extra, nargs| {
+            callback_error_ext(state, extra, true, |extra, nargs| {
                 // Lua ensures that `LUA_MINSTACK` stack spaces are available (after pushing arguments)
                 // The lock must be already held as the callback is executed
                 let rawlua = (*extra).raw_lua();
@@ -1251,7 +1262,7 @@ impl RawLua {
 
         unsafe extern "C-unwind" fn poll_future(state: *mut ffi::lua_State) -> c_int {
             let upvalue = get_userdata::<AsyncPollUpvalue>(state, ffi::lua_upvalueindex(1));
-            callback_error_ext(state, (*upvalue).extra.get(), |extra, _| {
+            callback_error_ext(state, (*upvalue).extra.get(), true, |extra, _| {
                 // Lua ensures that `LUA_MINSTACK` stack spaces are available (after pushing arguments)
                 // The lock must be already held as the future is polled
                 let rawlua = (*extra).raw_lua();
diff --git a/src/state/util.rs b/src/state/util.rs
@@ -27,6 +27,7 @@ impl Drop for StateGuard<'_> {
 pub(super) unsafe fn callback_error_ext<F, R>(
     state: *mut ffi::lua_State,
     mut extra: *mut ExtraData,
+    wrap_error: bool,
     f: F,
 ) -> R
 where
@@ -110,6 +111,13 @@ where
         Ok(Err(err)) => {
             let wrapped_error = prealloc_failure.r#use(state, extra);
 
+            if !wrap_error {
+                ptr::write(wrapped_error, WrappedFailure::Error(err));
+                get_internal_metatable::<WrappedFailure>(state);
+                ffi::lua_setmetatable(state, -2);
+                ffi::lua_error(state)
+            }
+
             // Build `CallbackError` with traceback
             let traceback = if ffi::lua_checkstack(state, ffi::LUA_TRACEBACK_STACK) != 0 {
                 ffi::luaL_traceback(state, state, ptr::null(), 0);
diff --git a/src/types.rs b/src/types.rs
@@ -90,6 +90,12 @@ pub(crate) type InterruptCallback = XRc<dyn Fn(&Lua) -> Result<VmState> + Send>;
 #[cfg(all(not(feature = "send"), feature = "luau"))]
 pub(crate) type InterruptCallback = XRc<dyn Fn(&Lua) -> Result<VmState>>;
 
+#[cfg(all(feature = "send", feature = "luau"))]
+pub(crate) type UserThreadCallback = XRc<dyn Fn(&Lua, crate::Value) -> Result<()> + Send>;
+
+#[cfg(all(not(feature = "send"), feature = "luau"))]
+pub(crate) type UserThreadCallback = XRc<dyn Fn(&Lua, crate::Value) -> Result<()>>;
+
 #[cfg(all(feature = "send", feature = "lua54"))]
 pub(crate) type WarnCallback = XRc<dyn Fn(&Lua, &str, bool) -> Result<()> + Send>;
 
diff --git a/tests/hooks.rs b/tests/hooks.rs
@@ -1,6 +1,5 @@
 #![cfg(not(feature = "luau"))]
 
-use std::ops::Deref;
 use std::sync::atomic::{AtomicI64, Ordering};
 use std::sync::{Arc, Mutex};
 
@@ -104,14 +103,10 @@ fn test_error_within_hook() -> Result<()> {
     })?;
 
     let err = lua.load("x = 1").exec().expect_err("panic didn't propagate");
-
     match err {
-        Error::CallbackError { cause, .. } => match cause.deref() {
-            Error::RuntimeError(s) => assert_eq!(s, "Something happened in there!"),
-            _ => panic!("wrong callback error kind caught"),
-        },
-        _ => panic!("wrong error kind caught"),
-    };
+        Error::RuntimeError(msg) => assert_eq!(msg, "Something happened in there!"),
+        err => panic!("expected `RuntimeError` with a specific message, got {err:?}"),
+    }
 
     Ok(())
 }
diff --git a/tests/luau.rs b/tests/luau.rs
@@ -2,8 +2,9 @@
 
 use std::fmt::Debug;
 use std::fs;
+use std::os::raw::c_void;
 use std::panic::{catch_unwind, AssertUnwindSafe};
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicPtr, AtomicU64, Ordering};
 use std::sync::Arc;
 
 use mlua::{Compiler, Error, Lua, LuaOptions, Result, StdLib, Table, ThreadStatus, Value, Vector, VmState};
@@ -395,11 +396,8 @@ fn test_interrupts() -> Result<()> {
     //
     lua.set_interrupt(|_| Err(Error::runtime("error from interrupt")));
     match f.call::<()>(()) {
-        Err(Error::CallbackError { cause, .. }) => match *cause {
-            Error::RuntimeError(ref m) if m == "error from interrupt" => {}
-            ref e => panic!("expected RuntimeError with a specific message, got {:?}", e),
-        },
-        r => panic!("expected CallbackError, got {:?}", r),
+        Err(Error::RuntimeError(ref msg)) => assert_eq!(msg, "error from interrupt"),
+        res => panic!("expected `RuntimeError` with a specific message, got {res:?}"),
     }
 
     lua.remove_interrupt();
@@ -412,3 +410,62 @@ fn test_fflags() {
     // We cannot really on any particular feature flag to be present
     assert!(Lua::set_fflag("UnknownFlag", true).is_err());
 }
+
+#[test]
+fn test_thread_events() -> Result<()> {
+    let lua = Lua::new();
+
+    let count = Arc::new(AtomicU64::new(0));
+    let thread_data: Arc<(AtomicPtr<c_void>, AtomicBool)> = Arc::new(Default::default());
+
+    let (count2, thread_data2) = (count.clone(), thread_data.clone());
+    lua.set_thread_event_callback(move |_, value| {
+        count2.fetch_add(1, Ordering::Relaxed);
+        (thread_data2.0).store(value.to_pointer() as *mut _, Ordering::Relaxed);
+        if value.is_thread() {
+            thread_data2.1.store(false, Ordering::Relaxed);
+        }
+        if value.is_light_userdata() {
+            thread_data2.1.store(true, Ordering::Relaxed);
+        }
+        Ok(())
+    });
+
+    let t = lua.create_thread(lua.load("return 123").into_function()?)?;
+    assert_eq!(count.load(Ordering::Relaxed), 1);
+    let t_ptr = t.to_pointer();
+    assert_eq!(t_ptr, thread_data.0.load(Ordering::Relaxed));
+    assert!(!thread_data.1.load(Ordering::Relaxed));
+
+    // Thead will be destroyed after GC cycle
+    drop(t);
+    lua.gc_collect()?;
+    assert_eq!(count.load(Ordering::Relaxed), 2);
+    assert_eq!(t_ptr, thread_data.0.load(Ordering::Relaxed));
+    assert!(thread_data.1.load(Ordering::Relaxed));
+
+    // Check that recursion is not allowed
+    let count3 = count.clone();
+    lua.set_thread_event_callback(move |lua, _value| {
+        count3.fetch_add(1, Ordering::Relaxed);
+        let _ = lua.create_thread(lua.load("return 123").into_function().unwrap())?;
+        Ok(())
+    });
+    let t = lua.create_thread(lua.load("return 123").into_function()?)?;
+    assert_eq!(count.load(Ordering::Relaxed), 3);
+
+    lua.remove_thread_event_callback();
+    drop(t);
+    lua.gc_collect()?;
+    assert_eq!(count.load(Ordering::Relaxed), 3);
+
+    // Test error inside callback
+    lua.set_thread_event_callback(move |_, _| Err(Error::runtime("error when processing thread event")));
+    let result = lua.create_thread(lua.load("return 123").into_function()?);
+    assert!(result.is_err());
+    assert!(
+        matches!(result, Err(Error::RuntimeError(err)) if err.contains("error when processing thread event"))
+    );
+
+    Ok(())
+}
diff --git a/tests/tests.rs b/tests/tests.rs
@@ -1308,8 +1308,7 @@ fn test_warnings() -> Result<()> {
     lua.set_warning_function(|_, _, _| Err(Error::runtime("warning error")));
     assert!(matches!(
         lua.load(r#"warn("test")"#).exec(),
-        Err(Error::CallbackError { cause, .. })
-            if matches!(*cause, Error::RuntimeError(ref err) if err == "warning error")
+        Err(Error::RuntimeError(ref err)) if err == "warning error"
     ));
 
     // Recursive warning
