Support thread create/resume/yield callbacks for all Lua/Luau versions

Introduce new `Lua::set_thread_event_callback` function to register callback for
thread events.
Thread collection callback is removed as unlikely very usefil and has too many limitations.

Author: khvzak

src/lib.rs

@@ -142,7 +142,7 @@ pub use crate::string::LuaString as String;
 #[doc(hidden)]
 pub use crate::table::{TablePairs, TableSequence};
 #[doc(hidden)]
-pub use crate::thread::ThreadStatus;
+pub use crate::thread::{ThreadEvent, ThreadStatus, ThreadTriggers};
 #[doc(hidden)]
 pub use crate::userdata::{
     MetaMethod, UserData, UserDataFields, UserDataMetatable, UserDataMethods, UserDataOwned, UserDataRef,

src/state.rs

@@ -22,7 +22,7 @@ use crate::scope::Scope;
 use crate::stdlib::StdLib;
 use crate::string::LuaString;
 use crate::table::Table;
-use crate::thread::Thread;
+use crate::thread::{Thread, ThreadEvent, ThreadTriggers};
 use crate::traits::{FromLua, FromLuaMulti, IntoLua, IntoLuaMulti};
 use crate::types::{
     AppDataRef, AppDataRefMut, ArcReentrantMutexGuard, Integer, LuaType, MaybeSend, MaybeSync, Number,
@@ -842,92 +842,87 @@ impl Lua {
         }
     }
 
-    /// 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)
+    /// Sets a callback invoked when thread lifecycle events occur.
+    ///
+    /// `triggers` controls which events trigger the callback, see [`ThreadTriggers`] for more
+    /// details.
+    ///
+    /// Only one callback can be registered at a time. Calling this again replaces the previous
+    /// callback and its triggers.
+    ///
+    /// # Example
+    ///
+    /// Subscribe only to yield events:
+    ///
+    /// ```
+    /// # use mlua::{Lua, Result, ThreadTriggers, ThreadEvent};
+    /// # fn main() -> Result<()> {
+    /// let lua = Lua::new();
+    /// lua.set_thread_event_callback(
+    ///     ThreadTriggers::ON_YIELD,
+    ///     |_lua, event| {
+    ///         if let ThreadEvent::Yield(thread) = event {
+    ///             println!("thread yielded");
+    ///         }
+    ///         Ok(())
+    ///     },
+    /// );
+    /// # Ok(())
+    /// # }
+    /// ```
+    pub fn set_thread_event_callback<F>(&self, triggers: ThreadTriggers, callback: F)
     where
-        F: Fn(&Lua, Thread) -> Result<()> + MaybeSend + 'static,
+        F: Fn(&Lua, ThreadEvent) -> 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);
+            (*lua.extra.get()).thread_triggers = triggers;
+            (*lua.extra.get()).thread_event_callback = Some(XRc::new(callback));
+            #[cfg(feature = "luau")]
+            {
+                let proc = Self::userthread_proc as _;
+                (*ffi::lua_callbacks(lua.main_state())).userthread = triggers.on_create.then(|| proc);
+            }
         }
     }
 
-    /// Sets a thread collection callback that will be called when a thread is destroyed.
+    /// Removes the thread event callback previously set by [`Lua::set_thread_event_callback`].
     ///
-    /// 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_collection_callback<F>(&self, callback: F)
-    where
-        F: Fn(crate::LightUserData) + MaybeSend + 'static,
-    {
+    /// This function has no effect if a callback was not previously set.
+    pub fn remove_thread_event_callback(&self) {
         let lua = self.lock();
+        let extra = lua.extra.get();
         unsafe {
-            (*lua.extra.get()).thread_collection_callback = Some(XRc::new(callback));
-            (*ffi::lua_callbacks(lua.main_state())).userthread = Some(Self::userthread_proc);
+            (*extra).thread_triggers = ThreadTriggers::new();
+            (*extra).thread_event_callback = None;
+            #[cfg(feature = "luau")]
+            {
+                (*ffi::lua_callbacks(lua.main_state())).userthread = None;
+            }
         }
     }
 
     #[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(&callback) > 2 {
-                return; // Don't allow recursion
-            }
-            ffi::lua_pushthread(child);
-            ffi::lua_xmove(child, (*extra).ref_thread, 1);
-            let value = Thread((*extra).raw_lua().pop_ref_thread(), child);
-            callback_error_ext(parent, extra, false, move |extra, _| {
-                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,
-            };
-
-            // 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;
+        // Only handle thread creation
+        if parent.is_null() {
+            return;
         }
-    }
 
-    /// 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_callbacks(&self) {
-        let lua = self.lock();
-        unsafe {
-            let extra = lua.extra.get();
-            (*extra).thread_creation_callback = None;
-            (*extra).thread_collection_callback = None;
-            (*ffi::lua_callbacks(lua.main_state())).userthread = None;
+        let extra = ExtraData::get(child);
+        if !(*extra).thread_triggers.on_create {
+            return;
         }
+        let callback = match &(*extra).thread_event_callback {
+            Some(cb) if XRc::strong_count(cb) == 1 => cb.clone(),
+            _ => return,
+        };
+        ffi::lua_pushthread(child);
+        ffi::lua_xmove(child, (*extra).ref_thread, 1);
+        let thread = Thread((*extra).raw_lua().pop_ref_thread(), child);
+        callback_error_ext(parent, extra, false, move |extra, _| {
+            callback((*extra).lua(), ThreadEvent::Create(thread))
+        })
     }
 
     /// Sets the warning function to be used by Lua to emit warnings.

src/state/extra.rs

@@ -12,7 +12,8 @@ use rustc_hash::FxHashMap;
 use crate::error::Result;
 use crate::state::RawLua;
 use crate::stdlib::StdLib;
-use crate::types::{AppData, ReentrantMutex, XRc};
+use crate::thread::ThreadTriggers;
+use crate::types::{AppData, ReentrantMutex, ThreadEventCallback, XRc};
 use crate::userdata::RawUserDataRegistry;
 use crate::util::{TypeKey, WrappedFailure, get_internal_metatable, push_internal_userdata};
 
@@ -81,10 +82,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) thread_creation_callback: Option<crate::types::ThreadCreationCallback>,
-    #[cfg(feature = "luau")]
-    pub(super) thread_collection_callback: Option<crate::types::ThreadCollectionCallback>,
+    pub(super) thread_triggers: ThreadTriggers,
+    pub(super) thread_event_callback: Option<ThreadEventCallback>,
 
     #[cfg(feature = "luau")]
     pub(crate) running_gc: bool,
@@ -186,10 +185,8 @@ impl ExtraData {
             warn_callback: None,
             #[cfg(feature = "luau")]
             interrupt_callback: None,
-            #[cfg(feature = "luau")]
-            thread_creation_callback: None,
-            #[cfg(feature = "luau")]
-            thread_collection_callback: None,
+            thread_triggers: ThreadTriggers::default(),
+            thread_event_callback: None,
             #[cfg(feature = "luau")]
             sandboxed: false,
             #[cfg(feature = "luau")]

src/state/raw.rs

@@ -15,11 +15,11 @@ use crate::state::util::callback_error_ext;
 use crate::stdlib::StdLib;
 use crate::string::LuaString;
 use crate::table::Table;
-use crate::thread::Thread;
+use crate::thread::{Thread, ThreadTriggers};
 use crate::traits::{FromLua, IntoLua};
 use crate::types::{
     AppDataRef, AppDataRefMut, Callback, CallbackUpvalue, DestructedUserdata, Integer, LightUserData,
-    LuaType, MaybeSend, ReentrantMutex, RegistryKey, ValueRef, XRc,
+    LuaType, MaybeSend, ReentrantMutex, RegistryKey, ThreadEventCallback, ValueRef, XRc,
 };
 use crate::userdata::{
     AnyUserData, MetaMethod, RawUserDataRegistry, UserData, UserDataRegistry, UserDataStorage,
@@ -643,7 +643,7 @@ impl RawLua {
 
         let protect = !self.unlikely_memory_error();
         #[cfg(feature = "luau")]
-        let protect = protect || (*self.extra.get()).thread_creation_callback.is_some();
+        let protect = protect || self.thread_event_triggers().on_create;
 
         let thread_state = if !protect {
             ffi::lua_newthread(state)
@@ -656,6 +656,19 @@ impl RawLua {
         self.set_thread_hook(thread_state, HookKind::Global)?;
 
         let thread = Thread(self.pop_ref(), thread_state);
+
+        // Exec creation callback for non-Luau (Luau handles this via `userthread_proc`)
+        #[cfg(not(feature = "luau"))]
+        if self.thread_event_triggers().on_create {
+            let extra = self.extra.get();
+            if let Some(ref cb) = (*extra).thread_event_callback
+                && XRc::strong_count(cb) == 1
+            {
+                let cb = cb.clone();
+                cb((*extra).lua(), crate::thread::ThreadEvent::Create(thread.clone()))?;
+            }
+        }
+
         ffi::lua_xpush(self.ref_thread(), thread_state, func.0.index);
         Ok(thread)
     }
@@ -691,6 +704,16 @@ impl RawLua {
         }
     }
 
+    #[inline(always)]
+    pub(crate) unsafe fn thread_event_triggers(&self) -> ThreadTriggers {
+        (*self.extra.get()).thread_triggers
+    }
+
+    #[inline(always)]
+    pub(crate) unsafe fn thread_event_callback(&self) -> Option<ThreadEventCallback> {
+        (*self.extra.get()).thread_event_callback.clone()
+    }
+
     /// Pushes a primitive type value onto the Lua stack.
     pub(crate) unsafe fn push_primitive_type<T: LuaType>(&self) -> bool {
         match T::TYPE_ID {