From 1b2b94c80899876b5fa6e0882f3437dfb81746a2 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Sat, 9 May 2020 00:18:48 +0100 Subject: [PATCH] Use Rust allocator for new Lua states that allows to set memory limit --- src/error.rs | 8 ++++ src/ffi/glue/glue.c | 4 ++ src/ffi/mod.rs | 22 +++++++++ src/lua.rs | 110 +++++++++++++++++++++++++++++++++++++++++++- tests/memory.rs | 31 +++++++++++++ 5 files changed, 173 insertions(+), 2 deletions(-) diff --git a/src/error.rs b/src/error.rs index 9c9d913..7813fd7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -36,6 +36,11 @@ pub enum Error { /// The Lua VM returns this error when there is an error running a `__gc` metamethod. #[cfg(any(feature = "lua53", feature = "lua52"))] GarbageCollectorError(StdString), + /// Setting memory limit is not available. + /// + /// This error can only happen when Lua state was not created by us and does not have the + /// custom allocator attached. + MemoryLimitNotAvailable, /// A mutable callback has triggered Lua code that has called the same mutable callback again. /// /// This is an error because a mutable callback can only be borrowed mutably once. @@ -148,6 +153,9 @@ impl fmt::Display for Error { Error::GarbageCollectorError(ref msg) => { write!(fmt, "garbage collector error: {}", msg) } + Error::MemoryLimitNotAvailable => { + write!(fmt, "setting memory limit is not available") + } Error::RecursiveMutCallback => write!(fmt, "mutable callback called recursively"), Error::CallbackDestructed => write!( fmt, diff --git a/src/ffi/glue/glue.c b/src/ffi/glue/glue.c index e8d74ab..e44fddb 100644 --- a/src/ffi/glue/glue.c +++ b/src/ffi/glue/glue.c @@ -84,6 +84,8 @@ const char *rs_int_type(int width) { return "i32"; case 8: return "i64"; + case 16: + return "i128"; } } @@ -96,6 +98,8 @@ const char *rs_uint_type(int width) { return "u32"; case 8: return "u64"; + case 16: + return "u128"; } } diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs index d46fa48..84718d5 100644 --- a/src/ffi/mod.rs +++ b/src/ffi/mod.rs @@ -253,6 +253,28 @@ pub use self::lualib::{LUA_FFILIBNAME, LUA_JITLIBNAME}; // Not actually defined in lua.h / luaconf.h pub const LUA_MAX_UPVALUES: c_int = 255; +// Copied from https://github.com/rust-lang/rust/blob/master/src/libstd/sys_common/alloc.rs +#[cfg(all(any( + target_arch = "x86", + target_arch = "arm", + target_arch = "mips", + target_arch = "powerpc", + target_arch = "powerpc64", + target_arch = "asmjs", + target_arch = "wasm32", + target_arch = "hexagon" +)))] +pub const SYS_MIN_ALIGN: usize = 8; +#[cfg(all(any( + target_arch = "x86_64", + target_arch = "aarch64", + target_arch = "mips64", + target_arch = "s390x", + target_arch = "sparc64", + target_arch = "riscv64" +)))] +pub const SYS_MIN_ALIGN: usize = 16; + #[allow(unused_imports, dead_code, non_camel_case_types)] mod glue { include!(concat!(env!("OUT_DIR"), "/glue.rs")); diff --git a/src/lua.rs b/src/lua.rs index 28e0a76..3bb9fbd 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -1,9 +1,10 @@ +use std::alloc; use std::any::TypeId; use std::cell::{RefCell, UnsafeCell}; use std::collections::HashMap; use std::ffi::CString; use std::marker::PhantomData; -use std::os::raw::{c_char, c_int}; +use std::os::raw::{c_char, c_int, c_void}; use std::sync::{Arc, Mutex}; use std::{mem, ptr, str}; @@ -51,12 +52,19 @@ struct ExtraData { registered_userdata: HashMap, registry_unref_list: Arc>>>, + mem_info: *mut MemoryInfo, + ref_thread: *mut ffi::lua_State, ref_stack_size: c_int, ref_stack_max: c_int, ref_free: Vec, } +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. @@ -91,6 +99,9 @@ impl Drop for Lua { ); *mlua_expect!(extra.registry_unref_list.lock(), "unref list poisoned") = None; ffi::lua_close(self.main_state); + if !extra.mem_info.is_null() { + Box::from_raw(extra.mem_info); + } } } } @@ -108,14 +119,74 @@ impl Lua { /// /// [`StdLib`]: struct.StdLib.html pub fn new_with(libs: StdLib) -> Lua { + unsafe extern "C" fn allocator( + extra_data: *mut c_void, + ptr: *mut c_void, + osize: usize, + nsize: usize, + ) -> *mut c_void { + let mem_info = &mut *(extra_data as *mut MemoryInfo); + + if nsize == 0 { + // Free memory + if !ptr.is_null() { + let layout = + alloc::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(); + } + + // 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(); + } + + let new_layout = alloc::Layout::from_size_align_unchecked(nsize, ffi::SYS_MIN_ALIGN); + + if ptr.is_null() { + // Allocate new memory + let new_ptr = alloc::alloc(new_layout) as *mut c_void; + if !new_ptr.is_null() { + mem_info.used_memory += mem_diff; + } + return new_ptr; + } + + // Reallocate memory + let old_layout = alloc::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() { + mem_info.used_memory += mem_diff; + } else if !ptr.is_null() && nsize < osize { + // Should not happend + alloc::handle_alloc_error(new_layout); + } + + new_ptr + } + unsafe { - let state = ffi::luaL_newstate(); + let mem_info = Box::into_raw(Box::new(MemoryInfo { + used_memory: 0, + memory_limit: 0, + })); + + let state = ffi::lua_newstate(allocator, mem_info as *mut c_void); ffi::luaL_requiref(state, cstr!("_G"), ffi::luaopen_base, 1); ffi::lua_pop(state, 1); let mut lua = Lua::init_from_ptr(state); lua.ephemeral = false; + lua.extra.lock().unwrap().mem_info = mem_info; mlua_expect!( protect_lua_closure(lua.main_state, 0, 0, |state| { @@ -203,6 +274,7 @@ impl Lua { registered_userdata: HashMap::new(), registry_unref_list: Arc::new(Mutex::new(Some(Vec::new()))), ref_thread, + mem_info: ptr::null_mut(), // We need 1 extra stack space to move values in and out of the ref stack. ref_stack_size: ffi::LUA_MINSTACK - 1, ref_stack_max: 0, @@ -237,6 +309,40 @@ impl Lua { unsafe { self.push_value(cb.call(())?).map(|_| 1) } } + /// Returns the amount of memory (in bytes) currently used inside this Lua state. + pub fn used_memory(&self) -> usize { + let extra = mlua_expect!(self.extra.lock(), "extra is poisoned"); + if extra.mem_info.is_null() { + // Get data from the Lua GC + unsafe { + 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); + return (used_kbytes as usize) * 1024 + (used_kbytes_rem as usize); + } + } + unsafe { (*extra.mem_info).used_memory as usize } + } + + /// Sets a memory limit on this Lua state. + /// + /// Once an allocation occurs that would pass this memory limit, + /// a `Error::MemoryError` is generated instead. + /// Returns previous limit (zero means no limit). + /// + /// Does not work on module mode where Lua state is managed externally. + #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))] + pub fn set_memory_limit(&self, memory_limit: usize) -> Result { + let mut extra = mlua_expect!(self.extra.lock(), "extra is poisoned"); + if extra.mem_info.is_null() { + return Err(Error::MemoryLimitNotAvailable); + } + unsafe { + let prev_limit = (*extra.mem_info).memory_limit as usize; + (*extra.mem_info).memory_limit = memory_limit as isize; + Ok(prev_limit) + } + } + /// Returns true if the garbage collector is currently running automatically. #[cfg(any(feature = "lua54", feature = "lua53", feature = "lua52"))] pub fn gc_is_running(&self) -> bool { diff --git a/tests/memory.rs b/tests/memory.rs index c53334b..aa0a809 100644 --- a/tests/memory.rs +++ b/tests/memory.rs @@ -2,6 +2,37 @@ use std::sync::Arc; use mlua::{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(); + + let initial_memory = lua.used_memory(); + assert!( + initial_memory > 0, + "used_memory reporting is wrong, lua uses memory for stdlib" + ); + + let f = lua + .load("local t = {}; for i = 1,10000 do t[i] = i end") + .into_function()?; + f.call::<_, ()>(()).expect("should trigger no memory limit"); + + lua.set_memory_limit(initial_memory + 10000)?; + match f.call::<_, ()>(()) { + Err(Error::MemoryError(_)) => {} + something_else => panic!("did not trigger memory error: {:?}", something_else), + }; + + lua.set_memory_limit(0)?; + f.call::<_, ()>(()).expect("should trigger no memory limit"); + + Ok(()) +} + #[test] fn test_gc_control() -> Result<()> { let lua = Lua::new();