Split Luau thread event callback to creation and collection callbacks.

We need this because they have different requirements:
- Thread creation callback can return Error or panic
- Thread collection callback runs during Luau GC cycle and cannot make any Lua calls or trigger panics.

Author: khvzak

src/state.rs

@@ -696,56 +696,91 @@ impl Lua {
         }
     }
 
-    /// Sets a thread event callback that will be called when a thread is created or destroyed.
+    /// Sets a thread creation callback that will be called when a thread is created.
+    #[cfg(any(feature = "luau", doc))]
+    #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
+    pub fn set_thread_creation_callback<F>(&self, callback: F)
+    where
+        F: Fn(&Lua, Thread) -> Result<()> + MaybeSend + 'static,
+    {
+        let lua = self.lock();
+        unsafe {
+            (*lua.extra.get()).thread_creation_callback = Some(XRc::new(callback));
+            (*ffi::lua_callbacks(lua.main_state())).userthread = Some(Self::userthread_proc);
+        }
+    }
+
+    /// Sets a thread collection callback that will be called when a thread is 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
+    /// Luau GC does not support exceptions during collection, so the callback must be
+    /// non-panicking. If the callback panics, the program will be aborted.
     #[cfg(any(feature = "luau", doc))]
     #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
-    pub fn set_thread_event_callback<F>(&self, callback: F)
+    pub fn set_thread_collection_callback<F>(&self, callback: F)
     where
-        F: Fn(&Lua, Value) -> Result<()> + MaybeSend + 'static,
+        F: Fn(crate::LightUserData) + 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 {
+        let lua = self.lock();
+        unsafe {
+            (*lua.extra.get()).thread_collection_callback = Some(XRc::new(callback));
+            (*ffi::lua_callbacks(lua.main_state())).userthread = Some(Self::userthread_proc);
+        }
+    }
+
+    #[cfg(feature = "luau")]
+    unsafe extern "C-unwind" fn userthread_proc(parent: *mut ffi::lua_State, child: *mut ffi::lua_State) {
+        let extra = ExtraData::get(child);
+        if !parent.is_null() {
+            // Thread is created
+            let callback = match (*extra).thread_creation_callback {
                 Some(ref cb) => cb.clone(),
                 None => return,
             };
-            if XRc::strong_count(&thread_cb) > 2 {
+            if XRc::strong_count(&callback) > 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))
-                }
-            };
+            ffi::lua_pushthread(child);
+            ffi::lua_xmove(child, (*extra).ref_thread, 1);
+            let value = Thread((*extra).raw_lua().pop_ref_thread(), child);
+            let _guard = StateGuard::new((*extra).raw_lua(), parent);
             callback_error_ext((*extra).raw_lua().state(), extra, false, move |extra, _| {
-                thread_cb((*extra).lua(), value)
+                callback((*extra).lua(), value)
             })
-        }
+        } else {
+            // Thread is about to be collected
+            let callback = match (*extra).thread_collection_callback {
+                Some(ref cb) => cb.clone(),
+                None => return,
+            };
 
-        // 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);
+            // We need to wrap the callback call in non-unwind function as it's not safe to unwind when
+            // Luau GC is running.
+            // This will trigger `abort()` if the callback panics.
+            unsafe extern "C" fn run_callback(
+                callback: *const crate::types::ThreadCollectionCallback,
+                value: *mut ffi::lua_State,
+            ) {
+                (*callback)(crate::LightUserData(value as _));
+            }
+
+            (*extra).running_gc = true;
+            run_callback(&callback, child);
+            (*extra).running_gc = false;
         }
     }
 
-    /// Removes any thread event callback previously set by `set_thread_event_callback`.
+    /// Removes any thread creation or collection callbacks previously set by
+    /// [`Lua::set_thread_creation_callback`] or [`Lua::set_thread_collection_callback`].
+    ///
+    /// This function has no effect if a thread callbacks were not previously set.
     #[cfg(any(feature = "luau", doc))]
     #[cfg_attr(docsrs, doc(cfg(feature = "luau")))]
-    pub fn remove_thread_event_callback(&self) {
+    pub fn remove_thread_callbacks(&self) {
         let lua = self.lock();
         unsafe {
-            (*lua.extra.get()).userthread_callback = None;
+            let extra = lua.extra.get();
+            (*extra).thread_creation_callback = None;
+            (*extra).thread_collection_callback = None;
             (*ffi::lua_callbacks(lua.main_state())).userthread = None;
         }
     }
@@ -2039,8 +2074,8 @@ impl Lua {
     pub(crate) fn lock(&self) -> ReentrantMutexGuard<RawLua> {
         let rawlua = self.raw.lock();
         #[cfg(feature = "luau")]
-        if unsafe { (*rawlua.extra.get()).running_userdata_gc } {
-            panic!("Luau VM is suspended while userdata destructor is running");
+        if unsafe { (*rawlua.extra.get()).running_gc } {
+            panic!("Luau VM is suspended while GC is running");
         }
         rawlua
     }
@@ -2066,8 +2101,8 @@ impl WeakLua {
     pub(crate) fn lock(&self) -> LuaGuard {
         let guard = LuaGuard::new(self.0.upgrade().expect("Lua instance is destroyed"));
         #[cfg(feature = "luau")]
-        if unsafe { (*guard.extra.get()).running_userdata_gc } {
-            panic!("Luau VM is suspended while userdata destructor is running");
+        if unsafe { (*guard.extra.get()).running_gc } {
+            panic!("Luau VM is suspended while GC is running");
         }
         guard
     }

src/state/extra.rs

@@ -81,10 +81,12 @@ pub(crate) struct ExtraData {
     #[cfg(feature = "luau")]
     pub(super) interrupt_callback: Option<crate::types::InterruptCallback>,
     #[cfg(feature = "luau")]
-    pub(super) userthread_callback: Option<crate::types::UserThreadCallback>,
+    pub(super) thread_creation_callback: Option<crate::types::ThreadCreationCallback>,
+    #[cfg(feature = "luau")]
+    pub(super) thread_collection_callback: Option<crate::types::ThreadCollectionCallback>,
 
     #[cfg(feature = "luau")]
-    pub(crate) running_userdata_gc: bool,
+    pub(crate) running_gc: bool,
     #[cfg(feature = "luau")]
     pub(super) sandboxed: bool,
     #[cfg(feature = "luau")]
@@ -181,15 +183,17 @@ impl ExtraData {
             #[cfg(feature = "luau")]
             interrupt_callback: None,
             #[cfg(feature = "luau")]
-            userthread_callback: None,
+            thread_creation_callback: None,
+            #[cfg(feature = "luau")]
+            thread_collection_callback: None,
             #[cfg(feature = "luau")]
             sandboxed: false,
             #[cfg(feature = "luau")]
             compiler: None,
             #[cfg(feature = "luau-jit")]
             enable_jit: true,
             #[cfg(feature = "luau")]
-            running_userdata_gc: false,
+            running_gc: false,
         }));
 
         // Store it in the registry

src/state/raw.rs

@@ -573,7 +573,7 @@ impl RawLua {
 
         let protect = !self.unlikely_memory_error();
         #[cfg(feature = "luau")]
-        let protect = protect || (*self.extra.get()).userthread_callback.is_some();
+        let protect = protect || (*self.extra.get()).thread_creation_callback.is_some();
 
         let thread_state = if !protect {
             ffi::lua_newthread(state)

src/types.rs

@@ -91,10 +91,16 @@ pub(crate) type InterruptCallback = XRc<dyn Fn(&Lua) -> Result<VmState> + Send>;
 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>;
+pub(crate) type ThreadCreationCallback = XRc<dyn Fn(&Lua, crate::Thread) -> Result<()> + Send>;
 
 #[cfg(all(not(feature = "send"), feature = "luau"))]
-pub(crate) type UserThreadCallback = XRc<dyn Fn(&Lua, crate::Value) -> Result<()>>;
+pub(crate) type ThreadCreationCallback = XRc<dyn Fn(&Lua, crate::Thread) -> Result<()>>;
+
+#[cfg(all(feature = "send", feature = "luau"))]
+pub(crate) type ThreadCollectionCallback = XRc<dyn Fn(crate::LightUserData) + Send>;
+
+#[cfg(all(not(feature = "send"), feature = "luau"))]
+pub(crate) type ThreadCollectionCallback = XRc<dyn Fn(crate::LightUserData)>;
 
 #[cfg(all(feature = "send", feature = "lua54"))]
 pub(crate) type WarnCallback = XRc<dyn Fn(&Lua, &str, bool) -> Result<()> + Send>;

src/userdata/util.rs

@@ -443,11 +443,11 @@ pub(crate) unsafe extern "C" fn collect_userdata<T>(
     // Almost none Lua operations are allowed when destructor is running,
     // so we need to set a flag to prevent calling any Lua functions
     let extra = (*ffi::lua_callbacks(state)).userdata as *mut crate::state::ExtraData;
-    (*extra).running_userdata_gc = true;
+    (*extra).running_gc = true;
     // Luau does not support _any_ panics in destructors (they are declared as "C", NOT as "C-unwind"),
     // so any panics will trigger `abort()`.
     ptr::drop_in_place(ud as *mut T);
-    (*extra).running_userdata_gc = false;
+    (*extra).running_gc = false;
 }
 
 // This method can be called by user or Lua GC to destroy the userdata.

tests/luau.rs

@@ -1,5 +1,6 @@
 #![cfg(feature = "luau")]
 
+use std::cell::Cell;
 use std::fmt::Debug;
 use std::fs;
 use std::os::raw::c_void;
@@ -419,17 +420,19 @@ fn test_thread_events() -> Result<()> {
     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| {
+    lua.set_thread_creation_callback(move |_, thread| {
         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);
-        }
+        (thread_data2.0).store(thread.to_pointer() as *mut _, Ordering::Relaxed);
+        thread_data2.1.store(false, Ordering::Relaxed);
         Ok(())
     });
+    let (count3, thread_data3) = (count.clone(), thread_data.clone());
+    lua.set_thread_collection_callback(move |thread_ptr| {
+        count3.fetch_add(1, Ordering::Relaxed);
+        if thread_data3.0.load(Ordering::Relaxed) == thread_ptr.0 {
+            thread_data3.1.store(true, Ordering::Relaxed);
+        }
+    });
 
     let t = lua.create_thread(lua.load("return 123").into_function()?)?;
     assert_eq!(count.load(Ordering::Relaxed), 1);
@@ -445,27 +448,47 @@ fn test_thread_events() -> Result<()> {
     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 count4 = count.clone();
+    lua.set_thread_creation_callback(move |lua, _value| {
+        count4.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();
+    lua.remove_thread_callbacks();
     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")));
+    lua.set_thread_creation_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"))
     );
 
+    // Test context switch when running Lua script
+    let count = Cell::new(0);
+    lua.set_thread_creation_callback(move |_, _| {
+        count.set(count.get() + 1);
+        if count.get() == 2 {
+            return Err(Error::runtime("thread limit exceeded"));
+        }
+        Ok(())
+    });
+    let result = lua
+        .load(
+            r#"
+            local co = coroutine.wrap(function() return coroutine.create(print) end)
+            co()
+    "#,
+        )
+        .exec();
+    assert!(result.is_err());
+    assert!(matches!(result, Err(Error::RuntimeError(err)) if err.contains("thread limit exceeded")));
+
     Ok(())
 }