Prepare for custom (Luau) userdata destructors.

khvzak · khvzak · commit 53ff494ab9fc · 2025-04-04T16:36:06.000+01:00
We need to add logic later to prevent calling any Lua functions when UserData destructor is running.
diff --git a/mlua-sys/Cargo.toml b/mlua-sys/Cargo.toml
@@ -40,7 +40,7 @@ cfg-if = "1.0"
 pkg-config = "0.3.17"
 lua-src = { version = ">= 547.0.0, < 547.1.0", optional = true }
 luajit-src = { version = ">= 210.5.0, < 210.6.0", optional = true }
-luau0-src = { version = "0.13.0", optional = true }
+luau0-src = { git = "https://github.com/mlua-rs/luau-src-rs", optional = true }
 
 [lints.rust]
 unexpected_cfgs = { level = "allow", check-cfg = ['cfg(raw_dylib)'] }
diff --git a/mlua-sys/src/luau/compat.rs b/mlua-sys/src/luau/compat.rs
@@ -372,7 +372,7 @@ pub unsafe fn luaL_loadbufferenv(
         fn free(p: *mut c_void);
     }
 
-    unsafe extern "C-unwind" fn data_dtor(data: *mut c_void) {
+    unsafe extern "C-unwind" fn data_dtor(_: *mut lua_State, data: *mut c_void) {
         free(*(data as *mut *mut c_char) as *mut c_void);
     }
 
diff --git a/mlua-sys/src/luau/lua.rs b/mlua-sys/src/luau/lua.rs
@@ -84,7 +84,6 @@ pub type lua_CFunction = unsafe extern "C-unwind" fn(L: *mut lua_State) -> c_int
 pub type lua_Continuation = unsafe extern "C-unwind" fn(L: *mut lua_State, status: c_int) -> c_int;
 
 /// Type for userdata destructor functions.
-pub type lua_Udestructor = unsafe extern "C-unwind" fn(*mut c_void);
 pub type lua_Destructor = unsafe extern "C-unwind" fn(L: *mut lua_State, *mut c_void);
 
 /// Type for memory-allocation functions.
@@ -191,7 +190,7 @@ extern "C-unwind" {
     pub fn lua_pushlightuserdatatagged(L: *mut lua_State, p: *mut c_void, tag: c_int);
     pub fn lua_newuserdatatagged(L: *mut lua_State, sz: usize, tag: c_int) -> *mut c_void;
     pub fn lua_newuserdatataggedwithmetatable(L: *mut lua_State, sz: usize, tag: c_int) -> *mut c_void;
-    pub fn lua_newuserdatadtor(L: *mut lua_State, sz: usize, dtor: lua_Udestructor) -> *mut c_void;
+    pub fn lua_newuserdatadtor(L: *mut lua_State, sz: usize, dtor: lua_Destructor) -> *mut c_void;
 
     pub fn lua_newbuffer(L: *mut lua_State, sz: usize) -> *mut c_void;
 
@@ -345,12 +344,14 @@ pub unsafe fn lua_newuserdata(L: *mut lua_State, sz: usize) -> *mut c_void {
 }
 
 #[inline(always)]
-pub unsafe fn lua_newuserdata_t<T>(L: *mut lua_State) -> *mut T {
-    unsafe extern "C-unwind" fn destructor<T>(ud: *mut c_void) {
+pub unsafe fn lua_newuserdata_t<T>(L: *mut lua_State, data: T) -> *mut T {
+    unsafe extern "C-unwind" fn destructor<T>(_: *mut lua_State, ud: *mut c_void) {
         ptr::drop_in_place(ud as *mut T);
     }
 
-    lua_newuserdatadtor(L, mem::size_of::<T>(), destructor::<T>) as *mut T
+    let ud_ptr = lua_newuserdatadtor(L, const { mem::size_of::<T>() }, destructor::<T>) as *mut T;
+    ptr::write(ud_ptr, data);
+    ud_ptr
 }
 
 // TODO: lua_strlen
diff --git a/src/luau/package.rs b/src/luau/package.rs
@@ -124,8 +124,7 @@ unsafe extern "C-unwind" fn lua_require(state: *mut ffi::lua_State) -> c_int {
     ffi::lua_pop(state, 1); // remove nil
 
     // load the module
-    let err_buf = ffi::lua_newuserdata_t::<StdString>(state);
-    err_buf.write(StdString::new());
+    let err_buf = ffi::lua_newuserdata_t(state, StdString::new());
     ffi::luaL_getsubtable(state, ffi::LUA_REGISTRYINDEX, cstr!("_LOADERS")); // _LOADERS is at index 3
     for i in 1.. {
         if ffi::lua_rawgeti(state, -1, i) == ffi::LUA_TNIL {
diff --git a/src/scope.rs b/src/scope.rs
@@ -168,7 +168,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
             #[cfg(feature = "luau")]
             let ud_ptr = {
                 let data = UserDataStorage::new_scoped(data);
-                util::push_userdata::<UserDataStorage<T>>(state, data, protect)?
+                util::push_userdata(state, data, protect)?
             };
             #[cfg(not(feature = "luau"))]
             let ud_ptr = util::push_uninit_userdata::<UserDataStorage<T>>(state, protect)?;
@@ -216,7 +216,7 @@ impl<'scope, 'env: 'scope> Scope<'scope, 'env> {
             #[cfg(feature = "luau")]
             let ud_ptr = {
                 let data = UserDataStorage::new_scoped(data);
-                util::push_userdata::<UserDataStorage<T>>(state, data, protect)?
+                util::push_userdata(state, data, protect)?
             };
             #[cfg(not(feature = "luau"))]
             let ud_ptr = util::push_uninit_userdata::<UserDataStorage<T>>(state, protect)?;
diff --git a/src/userdata.rs b/src/userdata.rs
@@ -30,7 +30,8 @@ pub use r#ref::{UserDataRef, UserDataRefMut};
 pub use registry::UserDataRegistry;
 pub(crate) use registry::{RawUserDataRegistry, UserDataProxy};
 pub(crate) use util::{
-    borrow_userdata_scoped, borrow_userdata_scoped_mut, init_userdata_metatable, TypeIdHints,
+    borrow_userdata_scoped, borrow_userdata_scoped_mut, collect_userdata, init_userdata_metatable,
+    TypeIdHints,
 };
 
 /// Kinds of metamethods that can be overridden.
diff --git a/src/userdata/registry.rs b/src/userdata/registry.rs
@@ -97,7 +97,7 @@ impl<T> UserDataRegistry<T> {
             meta_methods: Vec::new(),
             #[cfg(feature = "async")]
             async_meta_methods: Vec::new(),
-            destructor: super::util::userdata_destructor::<T>,
+            destructor: super::util::destroy_userdata_storage::<T>,
             type_id: r#type.type_id(),
             type_name: short_type_name::<T>(),
         };
diff --git a/src/userdata/util.rs b/src/userdata/util.rs
@@ -2,6 +2,7 @@ use std::any::TypeId;
 use std::cell::Cell;
 use std::marker::PhantomData;
 use std::os::raw::c_int;
+use std::ptr;
 
 use super::UserDataStorage;
 use crate::error::{Error, Result};
@@ -423,7 +424,29 @@ unsafe fn init_userdata_metatable_newindex(state: *mut ffi::lua_State) -> Result
     })
 }
 
-pub(super) unsafe extern "C-unwind" fn userdata_destructor<T>(state: *mut ffi::lua_State) -> c_int {
+// This method is called by Lua GC when it's time to collect the userdata.
+//
+// This method is usually used to collect internal userdata.
+#[cfg(not(feature = "luau"))]
+pub(crate) unsafe extern "C-unwind" fn collect_userdata<T>(state: *mut ffi::lua_State) -> c_int {
+    let ud = get_userdata::<T>(state, -1);
+    ptr::drop_in_place(ud);
+    0
+}
+
+// This method is called by Luau GC when it's time to collect the userdata.
+#[cfg(feature = "luau")]
+pub(crate) unsafe extern "C-unwind" fn collect_userdata<T>(
+    _state: *mut ffi::lua_State,
+    ud: *mut std::os::raw::c_void,
+) {
+    ptr::drop_in_place(ud as *mut T);
+}
+
+// This method can be called by user or Lua GC to destroy the userdata.
+// It checks if the userdata is safe to destroy and sets the "destroyed" metatable
+// to prevent further GC collection.
+pub(super) unsafe extern "C-unwind" fn destroy_userdata_storage<T>(state: *mut ffi::lua_State) -> c_int {
     let ud = get_userdata::<UserDataStorage<T>>(state, -1);
     if (*ud).is_safe_to_destroy() {
         take_userdata::<UserDataStorage<T>>(state);
diff --git a/src/util/error.rs b/src/util/error.rs
@@ -9,9 +9,8 @@ use std::sync::Arc;
 use crate::error::{Error, Result};
 use crate::memory::MemoryState;
 use crate::util::{
-    check_stack, get_internal_metatable, get_internal_userdata, init_internal_metatable,
-    push_internal_userdata, push_string, push_table, rawset_field, to_string, TypeKey,
-    DESTRUCTED_USERDATA_METATABLE,
+    check_stack, get_internal_userdata, init_internal_metatable, push_internal_userdata, push_string,
+    push_table, rawset_field, to_string, TypeKey, DESTRUCTED_USERDATA_METATABLE,
 };
 
 static WRAPPED_FAILURE_TYPE_KEY: u8 = 0;
@@ -31,12 +30,8 @@ impl TypeKey for WrappedFailure {
 
 impl WrappedFailure {
     pub(crate) unsafe fn new_userdata(state: *mut ffi::lua_State) -> *mut Self {
-        #[cfg(feature = "luau")]
-        let ud = ffi::lua_newuserdata_t::<Self>(state);
-        #[cfg(not(feature = "luau"))]
-        let ud = ffi::lua_newuserdata(state, std::mem::size_of::<Self>()) as *mut Self;
-        ptr::write(ud, WrappedFailure::None);
-        ud
+        // Unprotected calls always return `Ok`
+        push_internal_userdata(state, WrappedFailure::None, false).unwrap()
     }
 }
 
@@ -90,16 +85,11 @@ where
             let cause = Arc::new(err);
             let wrapped_error = WrappedFailure::Error(Error::CallbackError { traceback, cause });
             ptr::write(ud, wrapped_error);
-            get_internal_metatable::<WrappedFailure>(state);
-            ffi::lua_setmetatable(state, -2);
-
             ffi::lua_error(state)
         }
         Err(p) => {
             ffi::lua_settop(state, 1);
             ptr::write(ud, WrappedFailure::Panic(Some(p)));
-            get_internal_metatable::<WrappedFailure>(state);
-            ffi::lua_setmetatable(state, -2);
             ffi::lua_error(state)
         }
     }
diff --git a/src/util/mod.rs b/src/util/mod.rs
@@ -13,12 +13,12 @@ pub(crate) use short_names::short_type_name;
 pub(crate) use types::TypeKey;
 pub(crate) use userdata::{
     get_destructed_userdata_metatable, get_internal_metatable, get_internal_userdata, get_userdata,
-    init_internal_metatable, push_internal_userdata, take_userdata, DESTRUCTED_USERDATA_METATABLE,
+    init_internal_metatable, push_internal_userdata, push_userdata, take_userdata,
+    DESTRUCTED_USERDATA_METATABLE,
 };
 
 #[cfg(not(feature = "luau"))]
 pub(crate) use userdata::push_uninit_userdata;
-pub(crate) use userdata::push_userdata;
 
 // Checks that Lua has enough free stack space for future stack operations. On failure, this will
 // panic with an internal error message.
diff --git a/src/util/userdata.rs b/src/util/userdata.rs
@@ -1,7 +1,8 @@
 use std::os::raw::{c_int, c_void};
-use std::ptr;
+use std::{mem, ptr};
 
 use crate::error::Result;
+use crate::userdata::collect_userdata;
 use crate::util::{check_stack, get_metatable_ptr, push_table, rawset_field, TypeKey};
 
 // Pushes the userdata and attaches a metatable with __gc method.
@@ -10,11 +11,30 @@ pub(crate) unsafe fn push_internal_userdata<T: TypeKey>(
     state: *mut ffi::lua_State,
     t: T,
     protect: bool,
-) -> Result<()> {
-    push_userdata(state, t, protect)?;
+) -> Result<*mut T> {
+    #[cfg(not(feature = "luau"))]
+    let ud_ptr = if protect {
+        protect_lua!(state, 0, 1, move |state| {
+            let ud_ptr = ffi::lua_newuserdata(state, const { mem::size_of::<T>() }) as *mut T;
+            ptr::write(ud_ptr, t);
+            ud_ptr
+        })?
+    } else {
+        let ud_ptr = ffi::lua_newuserdata(state, const { mem::size_of::<T>() }) as *mut T;
+        ptr::write(ud_ptr, t);
+        ud_ptr
+    };
+
+    #[cfg(feature = "luau")]
+    let ud_ptr = if protect {
+        protect_lua!(state, 0, 1, move |state| ffi::lua_newuserdata_t::<T>(state, t))?
+    } else {
+        ffi::lua_newuserdata_t::<T>(state, t)
+    };
+
     get_internal_metatable::<T>(state);
     ffi::lua_setmetatable(state, -2);
-    Ok(())
+    Ok(ud_ptr)
 }
 
 #[track_caller]
@@ -35,12 +55,7 @@ pub(crate) unsafe fn init_internal_metatable<T: TypeKey>(
 
     #[cfg(not(feature = "luau"))]
     {
-        unsafe extern "C-unwind" fn userdata_destructor<T>(state: *mut ffi::lua_State) -> c_int {
-            take_userdata::<T>(state);
-            0
-        }
-
-        ffi::lua_pushcfunction(state, userdata_destructor::<T>);
+        ffi::lua_pushcfunction(state, collect_userdata::<T>);
         rawset_field(state, -2, "__gc")?;
     }
 
@@ -86,24 +101,34 @@ pub(crate) unsafe fn get_internal_userdata<T: TypeKey>(
 pub(crate) unsafe fn push_uninit_userdata<T>(state: *mut ffi::lua_State, protect: bool) -> Result<*mut T> {
     if protect {
         protect_lua!(state, 0, 1, |state| {
-            ffi::lua_newuserdata(state, std::mem::size_of::<T>()) as *mut T
+            ffi::lua_newuserdata(state, const { mem::size_of::<T>() }) as *mut T
         })
     } else {
-        Ok(ffi::lua_newuserdata(state, std::mem::size_of::<T>()) as *mut T)
+        Ok(ffi::lua_newuserdata(state, const { mem::size_of::<T>() }) as *mut T)
     }
 }
 
 // Internally uses 3 stack spaces, does not call checkstack.
 #[inline]
 pub(crate) unsafe fn push_userdata<T>(state: *mut ffi::lua_State, t: T, protect: bool) -> Result<*mut T> {
+    let size = const { mem::size_of::<T>() };
+
     #[cfg(not(feature = "luau"))]
-    let ud_ptr = push_uninit_userdata(state, protect)?;
+    let ud_ptr = if protect {
+        protect_lua!(state, 0, 1, move |state| ffi::lua_newuserdata(state, size))?
+    } else {
+        ffi::lua_newuserdata(state, size)
+    } as *mut T;
+
     #[cfg(feature = "luau")]
     let ud_ptr = if protect {
-        protect_lua!(state, 0, 1, |state| { ffi::lua_newuserdata_t::<T>(state) })?
+        protect_lua!(state, 0, 1, |state| {
+            ffi::lua_newuserdatadtor(state, size, collect_userdata::<T>)
+        })?
     } else {
-        ffi::lua_newuserdata_t::<T>(state)
-    };
+        ffi::lua_newuserdatadtor(state, size, collect_userdata::<T>)
+    } as *mut T;
+
     ptr::write(ud_ptr, t);
     Ok(ud_ptr)
 }
diff --git a/tests/userdata.rs b/tests/userdata.rs
@@ -264,7 +264,7 @@ fn test_gc_userdata() -> Result<()> {
     impl UserData for MyUserdata {
         fn add_methods<M: UserDataMethods<Self>>(methods: &mut M) {
             methods.add_method("access", |_, this, ()| {
-                assert!(this.id == 123);
+                assert_eq!(this.id, 123);
                 Ok(())
             });
         }

Original file line number	Diff line number	Diff line change
`@@ -372,7 +372,7 @@ pub unsafe fn luaL_loadbufferenv(`
`372`	`372`	`fn free(p: *mut c_void);`
`373`	`373`	`}`
`374`	`374`
`375`		`- unsafe extern "C-unwind" fn data_dtor(data: *mut c_void) {`
	`375`	`+ unsafe extern "C-unwind" fn data_dtor(_: mut lua_State, data: mut c_void) {`
`376`	`376`	`free((data as mut mut c_char) as mut c_void);`
`377`	`377`	`}`
`378`	`378`