Use Rust allocator for new Lua states that allows to set memory limit
This commit is contained in:
parent
526e7418d8
commit
1b2b94c808
|
@ -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,
|
||||
|
|
|
@ -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";
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -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"));
|
||||
|
|
110
src/lua.rs
110
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<TypeId, c_int>,
|
||||
registry_unref_list: Arc<Mutex<Option<Vec<c_int>>>>,
|
||||
|
||||
mem_info: *mut MemoryInfo,
|
||||
|
||||
ref_thread: *mut ffi::lua_State,
|
||||
ref_stack_size: c_int,
|
||||
ref_stack_max: c_int,
|
||||
ref_free: Vec<c_int>,
|
||||
}
|
||||
|
||||
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<usize> {
|
||||
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 {
|
||||
|
|
|
@ -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();
|
||||
|
|
Loading…
Reference in New Issue