Skip to content

Commit

Permalink
Remove MemoryState from ExtraData and retrieve using lua_getallocf (r…
Browse files Browse the repository at this point in the history 
…ecently added to Luau)
  • Loading branch information
@khvzak
khvzak committed Oct 23, 2023
1 parent f5021da commit a1e39a8
Show file tree
Hide file tree
Showing 4 changed files with 73 additions and 84 deletions.
2 changes: 1 addition & 1 deletion mlua-sys/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,4 +40,4 @@ cfg-if = "1.0"
pkg-config = "0.3.17"
lua-src = { version = ">= 546.0.0, < 546.1.0", optional = true }
luajit-src = { version = ">= 210.5.0, < 210.6.0", optional = true }
luau0-src = { version = "0.7.0", optional = true }
luau0-src = { version = "0.7.7", optional = true }
1 change: 1 addition & 0 deletions mlua-sys/src/luau/lua.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -278,6 +278,7 @@ extern "C-unwind" {
pub fn lua_getuserdatadtor(L: *mut lua_State, tag: c_int) -> Option<lua_Destructor>;
pub fn lua_clonefunction(L: *mut lua_State, idx: c_int);
pub fn lua_cleartable(L: *mut lua_State, idx: c_int);
pub fn lua_getallocf(L: *mut lua_State, ud: *mut *mut c_void) -> lua_Alloc;
}

//
Expand Down
102 changes: 48 additions & 54 deletions src/lua.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,14 @@ use std::cell::{RefCell, UnsafeCell};
use std::ffi::{CStr, CString};
use std::fmt;
use std::marker::PhantomData;
use std::mem::MaybeUninit;
use std::mem::{self, MaybeUninit};
use std::ops::Deref;
use std::os::raw::{c_char, c_int, c_void};
use std::panic::{catch_unwind, resume_unwind, AssertUnwindSafe, Location};
use std::ptr::NonNull;
use std::ptr;
use std::result::Result as StdResult;
use std::sync::atomic::{AtomicPtr, Ordering};
use std::sync::{Arc, Mutex};
use std::{mem, ptr, str};

use rustc_hash::FxHashMap;

Expand Down Expand Up @@ -60,6 +59,7 @@ use {
crate::types::{AsyncCallback, AsyncCallbackUpvalue, AsyncPollUpvalue},
futures_util::future::{self, Future},
futures_util::task::{noop_waker_ref, Context, Poll, Waker},
std::ptr::NonNull,
};

#[cfg(feature = "serialize")]
Expand Down Expand Up @@ -94,7 +94,6 @@ pub(crate) struct ExtraData {

safe: bool,
libs: StdLib,
mem_state: Option<NonNull<MemoryState>>,
#[cfg(feature = "module")]
skip_memory_check: bool,

Expand Down Expand Up @@ -223,6 +222,7 @@ impl LuaOptions {

#[cfg(feature = "async")]
pub(crate) static ASYNC_POLL_PENDING: u8 = 0;
#[cfg(not(feature = "luau"))]
pub(crate) static EXTRA_REGISTRY_KEY: u8 = 0;

const WRAPPED_FAILURE_POOL_SIZE: usize = 64;
Expand All @@ -244,11 +244,14 @@ impl Drop for Lua {
impl Drop for LuaInner {
fn drop(&mut self) {
unsafe {
#[cfg(feature = "luau")]
{
(*ffi::lua_callbacks(self.state())).userdata = ptr::null_mut();
}
let mem_state = MemoryState::get(self.main_state);

ffi::lua_close(self.main_state);

// Deallocate MemoryState
if !mem_state.is_null() {
drop(Box::from_raw(mem_state));
}
}
}
}
Expand All @@ -261,9 +264,6 @@ impl Drop for ExtraData {
}

*mlua_expect!(self.registry_unref_list.lock(), "unref list poisoned") = None;
if let Some(mem_state) = self.mem_state {
drop(unsafe { Box::from_raw(mem_state.as_ptr()) });
}
}
}

Expand Down Expand Up @@ -383,12 +383,11 @@ impl Lua {

/// Creates a new Lua state with required `libs` and `options`
unsafe fn inner_new(libs: StdLib, options: LuaOptions) -> Lua {
let mut mem_state: *mut MemoryState = Box::into_raw(Box::default());
let mem_state: *mut MemoryState = Box::into_raw(Box::default());
let mut state = ffi::lua_newstate(ALLOCATOR, mem_state as *mut c_void);
// If state is null then switch to Lua internal allocator
if state.is_null() {
drop(Box::from_raw(mem_state));
mem_state = ptr::null_mut();
state = ffi::luaL_newstate();
}
assert!(!state.is_null(), "Failed to instantiate Lua VM");
Expand All @@ -404,7 +403,6 @@ impl Lua {

let lua = Lua::init_from_ptr(state);
let extra = lua.extra.get();
(*extra).mem_state = NonNull::new(mem_state);

mlua_expect!(
load_from_std_lib(state, libs),
Expand Down Expand Up @@ -514,7 +512,6 @@ impl Lua {
app_data: AppData::default(),
safe: false,
libs: StdLib::NONE,
mem_state: None,
#[cfg(feature = "module")]
skip_memory_check: false,
ref_thread,
Expand Down Expand Up @@ -547,14 +544,8 @@ impl Lua {

// Store it in the registry
mlua_expect!(
(|state| {
push_gc_userdata(state, Arc::clone(&extra), true)?;
protect_lua!(state, 1, 0, fn(state) {
let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, extra_key);
})
})(main_state),
"Error while storing extra data",
set_extra_data(main_state, &extra),
"Error while storing extra data"
);

// Register `DestructedUserdata` type
Expand All @@ -572,13 +563,6 @@ impl Lua {
);
assert_stack(main_state, ffi::LUA_MINSTACK);

// Set Luau callbacks userdata to extra data
// We can use global callbacks userdata since we don't allow C modules in Luau
#[cfg(feature = "luau")]
{
(*ffi::lua_callbacks(main_state)).userdata = extra.get() as *mut c_void;
}

let inner = Arc::new(LuaInner {
state: AtomicPtr::new(state),
main_state,
Expand Down Expand Up @@ -1098,9 +1082,9 @@ impl Lua {
/// Returns the amount of memory (in bytes) currently used inside this Lua state.
pub fn used_memory(&self) -> usize {
unsafe {
match (*self.extra.get()).mem_state.map(|x| x.as_ref()) {
Some(mem_state) => mem_state.used_memory(),
None => {
match MemoryState::get(self.main_state) {
mem_state if !mem_state.is_null() => (*mem_state).used_memory(),
_ => {
// Get data from the Lua GC
let used_kbytes = ffi::lua_gc(self.main_state, ffi::LUA_GCCOUNT, 0);
let used_kbytes_rem = ffi::lua_gc(self.main_state, ffi::LUA_GCCOUNTB, 0);
Expand All @@ -1119,9 +1103,9 @@ impl Lua {
/// Does not work in module mode where Lua state is managed externally.
pub fn set_memory_limit(&self, limit: usize) -> Result<usize> {
unsafe {
match (*self.extra.get()).mem_state.map(|mut x| x.as_mut()) {
Some(mem_state) => Ok(mem_state.set_memory_limit(limit)),
None => Err(Error::MemoryLimitNotAvailable),
match MemoryState::get(self.main_state) {
mem_state if !mem_state.is_null() => Ok((*mem_state).set_memory_limit(limit)),
_ => Err(Error::MemoryLimitNotAvailable),
}
}
}
Expand Down Expand Up @@ -3169,16 +3153,13 @@ impl Lua {
#[inline]
pub(crate) unsafe fn unlikely_memory_error(&self) -> bool {
// MemoryInfo is empty in module mode so we cannot predict memory limits
(*self.extra.get())
.mem_state
.map(|x| x.as_ref().memory_limit() == 0)
.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;
})
match MemoryState::get(self.main_state) {
mem_state if !mem_state.is_null() => (*mem_state).memory_limit() == 0,
#[cfg(feature = "module")]
_ => (*self.extra.get()).skip_memory_check, // Check the special flag (only for module mode)
#[cfg(not(feature = "module"))]
_ => false,
}
}

#[cfg(feature = "unstable")]
Expand Down Expand Up @@ -3223,14 +3204,6 @@ impl LuaInner {
}
}

impl ExtraData {
#[cfg(feature = "luau")]
#[inline]
pub(crate) fn mem_state(&self) -> NonNull<MemoryState> {
self.mem_state.unwrap()
}
}

struct StateGuard<'a>(&'a LuaInner, *mut ffi::lua_State);

impl<'a> StateGuard<'a> {
Expand All @@ -3251,6 +3224,15 @@ unsafe fn extra_data(state: *mut ffi::lua_State) -> *mut ExtraData {
(*ffi::lua_callbacks(state)).userdata as *mut ExtraData
}

#[cfg(feature = "luau")]
unsafe fn set_extra_data(
state: *mut ffi::lua_State,
extra: &Arc<UnsafeCell<ExtraData>>,
) -> Result<()> {
(*ffi::lua_callbacks(state)).userdata = extra.get() as *mut _;
Ok(())
}

#[cfg(not(feature = "luau"))]
unsafe fn extra_data(state: *mut ffi::lua_State) -> *mut ExtraData {
let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
Expand All @@ -3265,6 +3247,18 @@ unsafe fn extra_data(state: *mut ffi::lua_State) -> *mut ExtraData {
(*extra_ptr).get()
}

#[cfg(not(feature = "luau"))]
unsafe fn set_extra_data(
state: *mut ffi::lua_State,
extra: &Arc<UnsafeCell<ExtraData>>,
) -> Result<()> {
push_gc_userdata(state, Arc::clone(extra), true)?;
protect_lua!(state, 1, 0, fn(state) {
let extra_key = &EXTRA_REGISTRY_KEY as *const u8 as *const c_void;
ffi::lua_rawsetp(state, ffi::LUA_REGISTRYINDEX, extra_key);
})
}

// Creates required entries in the metatable cache (see `util::METATABLE_CACHE`)
pub(crate) fn init_metatable_cache(cache: &mut FxHashMap<TypeId, u8>) {
cache.insert(TypeId::of::<Arc<UnsafeCell<ExtraData>>>(), 0);
Expand Down
52 changes: 23 additions & 29 deletions src/memory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,6 @@ use std::alloc::{self, Layout};
use std::os::raw::c_void;
use std::ptr;

#[cfg(feature = "luau")]
use crate::lua::ExtraData;

pub(crate) static ALLOCATOR: ffi::lua_Alloc = allocator;

#[derive(Default)]
Expand All @@ -20,6 +17,21 @@ pub(crate) struct MemoryState {
}

impl MemoryState {
#[inline]
pub(crate) unsafe fn get(state: *mut ffi::lua_State) -> *mut Self {
let mut mem_state = ptr::null_mut();
#[cfg(feature = "luau")]
{
ffi::lua_getallocf(state, &mut mem_state);
mlua_assert!(!mem_state.is_null(), "Luau state has no allocator userdata");
}
#[cfg(not(feature = "luau"))]
if ffi::lua_getallocf(state, &mut mem_state) != ALLOCATOR {
mem_state = ptr::null_mut();
}
mem_state as *mut MemoryState
}

#[inline]
pub(crate) fn used_memory(&self) -> usize {
self.used_memory as usize
Expand All @@ -37,36 +49,21 @@ impl MemoryState {
prev_limit as usize
}

// This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit
// This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit/luau
// to bypass the memory limit (if set).
#[cfg(any(feature = "lua51", feature = "luajit"))]
#[cfg(any(feature = "lua51", feature = "luajit", feature = "luau"))]
#[inline]
pub(crate) unsafe fn relax_limit_with(state: *mut ffi::lua_State, f: impl FnOnce()) {
let mut mem_state: *mut c_void = ptr::null_mut();
if ffi::lua_getallocf(state, &mut mem_state) == ALLOCATOR {
(*(mem_state as *mut MemoryState)).ignore_limit = true;
let mem_state = Self::get(state);
if !mem_state.is_null() {
(*mem_state).ignore_limit = true;
f();
(*(mem_state as *mut MemoryState)).ignore_limit = false;
(*mem_state).ignore_limit = false;
} else {
f();
}
}

// Same as the above but for Luau
// It does not have `lua_getallocf` function, so instead we use `lua_callbacks`
#[cfg(feature = "luau")]
#[inline]
pub(crate) unsafe fn relax_limit_with(state: *mut ffi::lua_State, f: impl FnOnce()) {
let extra = (*ffi::lua_callbacks(state)).userdata as *mut ExtraData;
if extra.is_null() {
return f();
}
let mem_state = (*extra).mem_state();
(*mem_state.as_ptr()).ignore_limit = true;
f();
(*mem_state.as_ptr()).ignore_limit = false;
}

// Does nothing apart from calling `f()`, we don't need to bypass any limits
#[cfg(any(feature = "lua52", feature = "lua53", feature = "lua54"))]
#[inline]
Expand All @@ -76,12 +73,9 @@ impl MemoryState {

// Returns `true` if the memory limit was reached on the last memory operation
#[cfg(feature = "luau")]
#[inline]
pub(crate) unsafe fn limit_reached(state: *mut ffi::lua_State) -> bool {
let extra = (*ffi::lua_callbacks(state)).userdata as *mut ExtraData;
if extra.is_null() {
return false;
}
(*(*extra).mem_state().as_ptr()).limit_reached
(*Self::get(state)).limit_reached
}
}

Expand Down

0 comments on commit a1e39a8

Please sign in to comment.