-
-
Notifications
You must be signed in to change notification settings - Fork 146
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support setting memory limit for Lua 5.1/JIT/Luau
Other versions already support this feature. Closes #119
- Loading branch information
Showing
6 changed files
with
209 additions
and
99 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -90,6 +90,7 @@ mod hook; | |
| mod lua; | ||
| #[cfg(feature = "luau")] | ||
| mod luau; | ||
| mod memory; | ||
| mod multi; | ||
| mod scope; | ||
| mod stdlib; | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } |
Oops, something went wrong.