michael/mlua
michael
/
mlua
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Support setting memory limit for Lua 5.1/JIT/Luau 

Other versions already support this feature.
Closes #119
This commit is contained in:
Alex Orlenko 2023-03-26 00:02:35 +00:00
parent 9c1669020b
commit d9aac08b81
No known key found for this signature in database
GPG Key ID: 4C150C250863B96D
6 changed files with 209 additions and 99 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
src/function.rs
View File

@ -5,6 +5,7 @@ use std::slice;
use crate::error::{Error, Result};
use crate::ffi;
use crate::memory::MemoryState;
use crate::types::LuaRef;
use crate::util::{
assert_stack, check_stack, error_traceback, pop_error, ptr_to_cstr_bytes, StackGuard,
@ -118,7 +119,7 @@ impl<'lua> Function<'lua> {
let _sg = StackGuard::new(state);
check_stack(state, nargs + 3)?;
ffi::lua_pushcfunction(state, error_traceback);
MemoryState::relax_limit_with(state, || ffi::lua_pushcfunction(state, error_traceback));
let stack_start = ffi::lua_gettop(state);
lua.push_ref(&self.0);
for arg in args.drain_all() {

1
src/lib.rs
View File

@ -90,6 +90,7 @@ mod hook;
mod lua;
#[cfg(feature = "luau")]
mod luau;
mod memory;
mod multi;
mod scope;
mod stdlib;

116
src/lua.rs
View File

@ -19,6 +19,7 @@ use crate::error::{Error, Result};
use crate::ffi;
use crate::function::Function;
use crate::hook::Debug;
use crate::memory::{MemoryState, ALLOCATOR};
use crate::scope::Scope;
use crate::stdlib::StdLib;
use crate::string::String;
@ -96,7 +97,7 @@ pub(crate) struct ExtraData {
safe: bool,
libs: StdLib,
mem_info: Option<NonNull<MemoryInfo>>,
mem_state: Option<NonNull<MemoryState>>,
ref_thread: *mut ffi::lua_State,
ref_stack_size: c_int,
@ -131,12 +132,6 @@ pub(crate) struct ExtraData {
compiler: Option<Compiler>,
}
#[derive(Default)]
struct MemoryInfo {
used_memory: isize,
memory_limit: isize,
}
/// Mode of the Lua garbage collector (GC).
///
/// In Lua 5.4 GC can work in two modes: incremental and generational.
@ -269,8 +264,8 @@ impl Drop for ExtraData {
};
*mlua_expect!(self.registry_unref_list.lock(), "unref list poisoned") = None;
if let Some(mem_info) = self.mem_info {
drop(unsafe { Box::from_raw(mem_info.as_ptr()) });
if let Some(mem_state) = self.mem_state {
drop(unsafe { Box::from_raw(mem_state.as_ptr()) });
}
}
}
@ -368,77 +363,19 @@ impl Lua {
/// Creates a new Lua state with required `libs` and `options`
unsafe fn inner_new(libs: StdLib, options: LuaOptions) -> Lua {
unsafe extern "C" fn allocator(
extra_data: *mut c_void,
ptr: *mut c_void,
osize: usize,
nsize: usize,
) -> *mut c_void {
use std::alloc::{self, Layout};
let mem_info = &mut *(extra_data as *mut MemoryInfo);
if nsize == 0 {
// Free memory
if !ptr.is_null() {
let layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN);
alloc::dealloc(ptr as *mut u8, layout);
mem_info.used_memory -= osize as isize;
}
return ptr::null_mut();
}
// Do not allocate more than isize::MAX
if nsize > isize::MAX as usize {
return ptr::null_mut();
}
// Are we fit to the memory limits?
let mut mem_diff = nsize as isize;
if !ptr.is_null() {
mem_diff -= osize as isize;
}
let new_used_memory = mem_info.used_memory + mem_diff;
if mem_info.memory_limit > 0 && new_used_memory > mem_info.memory_limit {
return ptr::null_mut();
}
mem_info.used_memory += mem_diff;
if ptr.is_null() {
// Allocate new memory
let new_layout = match Layout::from_size_align(nsize, ffi::SYS_MIN_ALIGN) {
Ok(layout) => layout,
Err(_) => return ptr::null_mut(),
};
let new_ptr = alloc::alloc(new_layout) as *mut c_void;
if new_ptr.is_null() {
alloc::handle_alloc_error(new_layout);
}
return new_ptr;
}
// Reallocate memory
let old_layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN);
let new_ptr = alloc::realloc(ptr as *mut u8, old_layout, nsize) as *mut c_void;
if new_ptr.is_null() {
alloc::handle_alloc_error(old_layout);
}
new_ptr
}
// Skip Rust allocator for non-vendored LuaJIT (see https://github.com/khvzak/mlua/issues/176)
let use_rust_allocator = !(cfg!(feature = "luajit") && cfg!(not(feature = "vendored")));
let (state, mem_info) = if use_rust_allocator {
let mut mem_info: *mut MemoryInfo = Box::into_raw(Box::default());
let mut state = ffi::lua_newstate(allocator, mem_info as *mut c_void);
let (state, mem_state) = if use_rust_allocator {
let mut 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 (it's possible for LuaJIT on non-x86 arch) then switch to Lua internal allocator
if state.is_null() {
drop(Box::from_raw(mem_info));
mem_info = ptr::null_mut();
drop(Box::from_raw(mem_state));
mem_state = ptr::null_mut();
state = ffi::luaL_newstate();
}
(state, mem_info)
(state, mem_state)
} else {
(ffi::luaL_newstate(), ptr::null_mut())
};
@ -449,7 +386,7 @@ impl Lua {
let lua = Lua::init_from_ptr(state);
let extra = lua.extra.get();
(*extra).mem_info = NonNull::new(mem_info);
(*extra).mem_state = NonNull::new(mem_state);
mlua_expect!(
load_from_std_lib(state, libs),
@ -559,7 +496,7 @@ impl Lua {
app_data: RefCell::new(FxHashMap::default()),
safe: false,
libs: StdLib::NONE,
mem_info: None,
mem_state: None,
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,
@ -1124,8 +1061,8 @@ 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_info.map(|x| x.as_ref()) {
Some(mem_info) => mem_info.used_memory as usize,
match (*self.extra.get()).mem_state.map(|x| x.as_ref()) {
Some(mem_state) => mem_state.used_memory(),
None => {
// Get data from the Lua GC
let used_kbytes = ffi::lua_gc(self.main_state, ffi::LUA_GCCOUNT, 0);
@ -1143,17 +1080,10 @@ impl Lua {
/// Returns previous limit (zero means no limit).
///
/// Does not work on module mode where Lua state is managed externally.
///
/// Requires `feature = "lua54/lua53/lua52"`
#[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
pub fn set_memory_limit(&self, memory_limit: usize) -> Result<usize> {
pub fn set_memory_limit(&self, limit: usize) -> Result<usize> {
unsafe {
match (*self.extra.get()).mem_info.map(|mut x| x.as_mut()) {
Some(mem_info) => {
let prev_limit = mem_info.memory_limit as usize;
mem_info.memory_limit = memory_limit as isize;
Ok(prev_limit)
}
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),
}
}
@ -3022,8 +2952,8 @@ impl Lua {
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_info
.map(|x| x.as_ref().memory_limit == 0)
.mem_state
.map(|x| x.as_ref().memory_limit() == 0)
.unwrap_or_default()
}
@ -3063,6 +2993,14 @@ 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> {

153
src/memory.rs Normal file
View File

@ -0,0 +1,153 @@
use std::alloc::{self, Layout};
use std::os::raw::c_void;
use std::ptr;
use crate::ffi;
#[cfg(feature = "luau")]
use crate::lua::ExtraData;
pub(crate) static ALLOCATOR: ffi::lua_Alloc = allocator;
#[derive(Default)]
pub(crate) struct MemoryState {
used_memory: isize,
memory_limit: isize,
// Can be set to temporary ignore the memory limit.
// This is used when calling `lua_pushcfunction` for lua5.1/jit/luau.
ignore_limit: bool,
// Indicates that the memory limit was reached on the last allocation.
#[cfg(feature = "luau")]
limit_reached: bool,
}
impl MemoryState {
#[inline]
pub(crate) fn used_memory(&self) -> usize {
self.used_memory as usize
}
#[inline]
pub(crate) fn memory_limit(&self) -> usize {
self.memory_limit as usize
}
#[inline]
pub(crate) fn set_memory_limit(&mut self, limit: usize) -> usize {
let prev_limit = self.memory_limit;
self.memory_limit = limit as isize;
prev_limit as usize
}
// This function is used primarily for calling `lua_pushcfunction` in lua5.1/jit
// to bypass the memory limit (if set).
#[cfg(any(feature = "lua51", feature = "luajit"))]
#[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;
f();
(*(mem_state as *mut MemoryState)).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]
pub(crate) unsafe fn relax_limit_with(_state: *mut ffi::lua_State, f: impl FnOnce()) {
f();
}
// Returns `true` if the memory limit was reached on the last memory operation
#[cfg(feature = "luau")]
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
}
}
unsafe extern "C" fn allocator(
extra: *mut c_void,
ptr: *mut c_void,
osize: usize,
nsize: usize,
) -> *mut c_void {
let mem_state = &mut *(extra as *mut MemoryState);
#[cfg(feature = "luau")]
{
// Reset the flag
mem_state.limit_reached = false;
}
if nsize == 0 {
// Free memory
if !ptr.is_null() {
let layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN);
alloc::dealloc(ptr as *mut u8, layout);
mem_state.used_memory -= osize as isize;
}
return ptr::null_mut();
}
// Do not allocate more than isize::MAX
if nsize > isize::MAX as usize {
return ptr::null_mut();
}
// Are we fit to the memory limits?
let mut mem_diff = nsize as isize;
if !ptr.is_null() {
mem_diff -= osize as isize;
}
let mem_limit = mem_state.memory_limit;
let new_used_memory = mem_state.used_memory + mem_diff;
if mem_limit > 0 && new_used_memory > mem_limit && !mem_state.ignore_limit {
#[cfg(feature = "luau")]
{
mem_state.limit_reached = true;
}
return ptr::null_mut();
}
mem_state.used_memory += mem_diff;
if ptr.is_null() {
// Allocate new memory
let new_layout = match Layout::from_size_align(nsize, ffi::SYS_MIN_ALIGN) {
Ok(layout) => layout,
Err(_) => return ptr::null_mut(),
};
let new_ptr = alloc::alloc(new_layout) as *mut c_void;
if new_ptr.is_null() {
alloc::handle_alloc_error(new_layout);
}
return new_ptr;
}
// Reallocate memory
let old_layout = Layout::from_size_align_unchecked(osize, ffi::SYS_MIN_ALIGN);
let new_ptr = alloc::realloc(ptr as *mut u8, old_layout, nsize) as *mut c_void;
if new_ptr.is_null() {
alloc::handle_alloc_error(old_layout);
}
new_ptr
}

20
src/util.rs
View File

@ -12,6 +12,7 @@ use rustc_hash::FxHashMap;
use crate::error::{Error, Result};
use crate::ffi;
use crate::memory::MemoryState;
static METATABLE_CACHE: Lazy<FxHashMap<TypeId, u8>> = Lazy::new(|| {
let mut map = FxHashMap::with_capacity_and_hasher(32, Default::default());
@ -89,8 +90,10 @@ pub unsafe fn protect_lua_call(
) -> Result<()> {
let stack_start = ffi::lua_gettop(state) - nargs;
ffi::lua_pushcfunction(state, error_traceback);
ffi::lua_pushcfunction(state, f);
MemoryState::relax_limit_with(state, || {
ffi::lua_pushcfunction(state, error_traceback);
ffi::lua_pushcfunction(state, f);
});
if nargs > 0 {
ffi::lua_rotate(state, stack_start + 1, 2);
}
@ -147,8 +150,10 @@ where
let stack_start = ffi::lua_gettop(state) - nargs;
ffi::lua_pushcfunction(state, error_traceback);
ffi::lua_pushcfunction(state, do_call::<F, R>);
MemoryState::relax_limit_with(state, || {
ffi::lua_pushcfunction(state, error_traceback);
ffi::lua_pushcfunction(state, do_call::<F, R>);
});
if nargs > 0 {
ffi::lua_rotate(state, stack_start + 1, 2);
}
@ -662,6 +667,13 @@ where
}
pub unsafe extern "C" fn error_traceback(state: *mut ffi::lua_State) -> c_int {
// This is a workaround for bug in Luau, when it calls error handler for memory allocation error
// See https://github.com/Roblox/luau/issues/880
#[cfg(feature = "luau")]
if MemoryState::limit_reached(state) {
return 0;
}
if ffi::lua_checkstack(state, 2) == 0 {
// If we don't have enough stack space to even check the error type, do
// nothing so we don't risk shadowing a rust panic.

15
tests/memory.rs
View File

@ -1,11 +1,7 @@
use std::sync::Arc;
use mlua::{GCMode, Lua, Result, UserData};
use mlua::{Error, GCMode, Lua, Result, UserData};
#[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
use mlua::Error;
#[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))]
#[test]
fn test_memory_limit() -> Result<()> {
let lua = Lua::new();
@ -21,6 +17,15 @@ fn test_memory_limit() -> Result<()> {
.into_function()?;
f.call::<_, ()>(()).expect("should trigger no memory limit");
if cfg!(feature = "luajit") && cfg!(not(feature = "vendored")) {
// we don't support setting memory limit for non-vendored luajit
assert!(matches!(
lua.set_memory_limit(0),
Err(Error::MemoryLimitNotAvailable)
));
return Ok(());
}
lua.set_memory_limit(initial_memory + 10000)?;
match f.call::<_, ()>(()) {
Err(Error::MemoryError(_)) => {}