diff --git a/mlua_derive/src/lib.rs b/mlua_derive/src/lib.rs index 8f057dc..59f77a9 100644 --- a/mlua_derive/src/lib.rs +++ b/mlua_derive/src/lib.rs @@ -13,6 +13,7 @@ use { #[derive(Default)] struct ModuleAttributes { name: Option, + skip_memory_check: bool, } impl ModuleAttributes { @@ -26,6 +27,11 @@ impl ModuleAttributes { return Err(meta.error("`name` attribute must have a value")); } } + } else if meta.path.is_ident("skip_memory_check") { + if meta.value().is_ok() { + return Err(meta.error("`skip_memory_check` attribute have no values")); + } + self.skip_memory_check = true; } else { return Err(meta.error("unsupported module attribute")); } @@ -45,6 +51,7 @@ pub fn lua_module(attr: TokenStream, item: TokenStream) -> TokenStream { let func_name = &func.sig.ident; let module_name = args.name.unwrap_or_else(|| func_name.clone()); let ext_entrypoint_name = Ident::new(&format!("luaopen_{module_name}"), Span::call_site()); + let skip_memory_check = args.skip_memory_check; let wrapped = quote! { ::mlua::require_module_feature!(); @@ -53,7 +60,11 @@ pub fn lua_module(attr: TokenStream, item: TokenStream) -> TokenStream { #[no_mangle] unsafe extern "C" fn #ext_entrypoint_name(state: *mut ::mlua::lua_State) -> ::std::os::raw::c_int { - ::mlua::Lua::init_from_ptr(state) + let lua = ::mlua::Lua::init_from_ptr(state); + if #skip_memory_check { + lua.skip_memory_check(true); + } + lua .entrypoint1(#func_name) .expect("cannot initialize module") } diff --git a/src/lib.rs b/src/lib.rs index 8685128..16def39 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -235,7 +235,7 @@ pub use mlua_derive::chunk; /// /// You can also pass options to the attribute: /// -/// name - name of the module, defaults to the name of the function +/// * name - name of the module, defaults to the name of the function /// /// ```ignore /// #[mlua::lua_module(name = "alt_module")] @@ -244,6 +244,20 @@ pub use mlua_derive::chunk; /// } /// ``` /// +/// * skip_memory_check - skip memory allocation checks for some operations. +/// +/// In module mode, mlua runs in unknown environment and cannot say are there any memory +/// limits or not. As result, some operations that require memory allocation runs in +/// protected mode. Setting this mode will improve performance of such operations +/// with risk of having uncaught exceptions and memory leaks. +/// +/// ```ignore +/// #[mlua::lua_module(skip_memory_check)] +/// fn my_module(lua: &Lua) -> Result { +/// ... +/// } +/// ``` +/// #[cfg(any(feature = "module", docsrs))] #[cfg_attr(docsrs, doc(cfg(feature = "module")))] pub use mlua_derive::lua_module; diff --git a/src/lua.rs b/src/lua.rs index 0d356ba..5c80acf 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -95,6 +95,8 @@ pub(crate) struct ExtraData { safe: bool, libs: StdLib, mem_state: Option>, + #[cfg(feature = "module")] + skip_memory_check: bool, // Auxiliary thread to store references ref_thread: *mut ffi::lua_State, @@ -515,6 +517,8 @@ impl Lua { safe: false, libs: StdLib::NONE, mem_state: None, + #[cfg(feature = "module")] + skip_memory_check: false, ref_thread, // We need 1 extra stack space to move values in and out of the ref stack. ref_stack_size: ffi::LUA_MINSTACK - 1, @@ -773,6 +777,13 @@ impl Lua { self.entrypoint(move |lua, _: ()| func(lua)) } + /// Skips memory checks for some operations. + #[doc(hidden)] + #[cfg(feature = "module")] + pub fn skip_memory_check(&self, skip: bool) { + unsafe { (*self.extra.get()).skip_memory_check = skip }; + } + /// Enables (or disables) sandbox mode on this Lua instance. /// /// This method, in particular: @@ -3079,7 +3090,13 @@ impl Lua { (*self.extra.get()) .mem_state .map(|x| x.as_ref().memory_limit() == 0) - .unwrap_or_default() + .unwrap_or_else(|| { + // Alternatively, check the special flag (only for module mode) + #[cfg(feature = "module")] + return (*self.extra.get()).skip_memory_check; + #[cfg(not(feature = "module"))] + return false; + }) } #[cfg(feature = "unstable")] diff --git a/tests/module/src/lib.rs b/tests/module/src/lib.rs index 28b5b45..25689fe 100644 --- a/tests/module/src/lib.rs +++ b/tests/module/src/lib.rs @@ -26,7 +26,7 @@ struct MyUserData(i32); impl LuaUserData for MyUserData {} -#[mlua::lua_module(name = "test_module_second")] +#[mlua::lua_module(name = "test_module_second", skip_memory_check)] fn test_module2(lua: &Lua) -> LuaResult { let exports = lua.create_table()?; exports.set("userdata", MyUserData(123))?;