diff --git a/src/error.rs b/src/error.rs index 2a86eacc..f4a339ce 100644 --- a/src/error.rs +++ b/src/error.rs @@ -42,11 +42,11 @@ pub enum Error { GarbageCollectorError(StdString), /// Potentially unsafe action in safe mode. SafetyError(StdString), - /// Setting memory limit is not available. + /// Setting memory limit or getting/resetting the number of allocation 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, + MemoryStatsNotAvailable, /// 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. @@ -218,8 +218,8 @@ impl fmt::Display for Error { Error::SafetyError(ref msg) => { write!(fmt, "safety error: {msg}") }, - Error::MemoryLimitNotAvailable => { - write!(fmt, "setting memory limit is not available") + Error::MemoryStatsNotAvailable => { + write!(fmt, "memory stats information is not available") } Error::RecursiveMutCallback => write!(fmt, "mutable callback called recursively"), Error::CallbackDestructed => write!( diff --git a/src/lua.rs b/src/lua.rs index 11765cee..012b14bd 100644 --- a/src/lua.rs +++ b/src/lua.rs @@ -1121,7 +1121,37 @@ impl Lua { 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), + None => Err(Error::MemoryStatsNotAvailable), + } + } + } + + /// Returns the count of memory allocations invoked by this Lua runtime. + /// + /// Note that in case of overflow, the counter resets to zero(simply speaking, it 'wraps'), + /// and the overflowing threshold above the maximum depends on the platform. + /// + /// Returns [`Error::MemoryStatsNotAvailable`] if Lua state is managed externally. + pub fn num_allocations(&self) -> Result { + unsafe { + match (*self.extra.get()).mem_state.map(|x| x.as_ref()) { + Some(mem_state) => Ok(mem_state.num_allocations()), + None => Err(Error::MemoryStatsNotAvailable), + } + } + } + + /// Resets the count of memory allocations to zero. + /// + /// Returns [`Error::MemoryStatsNotAvailable`] if Lua state is managed externally. + pub fn reset_num_allocations(&self) -> Result<()> { + unsafe { + match (*self.extra.get()).mem_state.map(|mut x| x.as_mut()) { + Some(mem_state) => { + mem_state.reset_num_allocations(); + Ok(()) + } + None => Err(Error::MemoryStatsNotAvailable), } } } diff --git a/src/memory.rs b/src/memory.rs index e199a759..d489f24b 100644 --- a/src/memory.rs +++ b/src/memory.rs @@ -17,6 +17,7 @@ pub(crate) struct MemoryState { // Indicates that the memory limit was reached on the last allocation. #[cfg(feature = "luau")] limit_reached: bool, + num_allocations: usize, } impl MemoryState { @@ -37,6 +38,16 @@ impl MemoryState { prev_limit as usize } + #[inline] + pub(crate) fn num_allocations(&self) -> usize { + self.num_allocations + } + + #[inline] + pub(crate) fn reset_num_allocations(&mut self) { + self.num_allocations = 0; + } + // 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"))] @@ -139,6 +150,7 @@ unsafe extern "C-unwind" fn allocator( if new_ptr.is_null() { alloc::handle_alloc_error(new_layout); } + mem_state.num_allocations = mem_state.num_allocations.wrapping_add(1); return new_ptr; }