From f16aca687dbdd9e581ce10200f86889614e8eeb5 Mon Sep 17 00:00:00 2001 From: Alex Orlenko Date: Fri, 9 May 2025 12:48:40 +0100 Subject: [PATCH] Terminate underlying Rust future when `AsyncThread` is dropped. Before this change, Lua GC was responsible to collect and destroy the future if `AsyncThread` dropped in yielded state. Now we will propagate "drop" event immediately so Lua GC need to only free the memory. --- src/state.rs | 7 +++++ src/state/raw.rs | 64 +++++++++++++-------------------------- src/thread.rs | 78 +++++++++++++++++++++++++++++++++--------------- src/types.rs | 2 +- tests/async.rs | 14 ++++++++- 5 files changed, 95 insertions(+), 70 deletions(-) diff --git a/src/state.rs b/src/state.rs index 53a19e6..4be9ba4 100644 --- a/src/state.rs +++ b/src/state.rs @@ -2036,6 +2036,13 @@ impl Lua { LightUserData(&ASYNC_POLL_PENDING as *const u8 as *mut std::os::raw::c_void) } + #[cfg(feature = "async")] + #[inline(always)] + pub(crate) fn poll_terminate() -> LightUserData { + static ASYNC_POLL_TERMINATE: u8 = 0; + LightUserData(&ASYNC_POLL_TERMINATE as *const u8 as *mut std::os::raw::c_void) + } + /// Returns a weak reference to the Lua instance. /// /// This is useful for creating a reference to the Lua instance that does not prevent it from diff --git a/src/state/raw.rs b/src/state/raw.rs index 44439da..cc4ce1c 100644 --- a/src/state/raw.rs +++ b/src/state/raw.rs @@ -613,46 +613,11 @@ impl RawLua { self.create_thread(func) } - /// Resets thread (coroutine) and returns it to the pool for later use. + /// Returns the thread to the pool for later use. #[cfg(feature = "async")] pub(crate) unsafe fn recycle_thread(&self, thread: &mut Thread) { - let thread_state = thread.1; let extra = &mut *self.extra.get(); - if extra.thread_pool.len() == extra.thread_pool.capacity() { - #[cfg(feature = "lua54")] - if ffi::lua_status(thread_state) != ffi::LUA_OK { - // Close all to-be-closed variables without returning thread to the pool - #[cfg(not(feature = "vendored"))] - ffi::lua_resetthread(thread_state); - #[cfg(feature = "vendored")] - ffi::lua_closethread(thread_state, self.state()); - } - return; - } - - let mut reset_ok = false; - if ffi::lua_status(thread_state) == ffi::LUA_OK { - if ffi::lua_gettop(thread_state) > 0 { - ffi::lua_settop(thread_state, 0); - } - reset_ok = true; - } - - #[cfg(feature = "lua54")] - if !reset_ok { - #[cfg(not(feature = "vendored"))] - let status = ffi::lua_resetthread(thread_state); - #[cfg(feature = "vendored")] - let status = ffi::lua_closethread(thread_state, self.state()); - reset_ok = status == ffi::LUA_OK; - } - #[cfg(feature = "luau")] - if !reset_ok { - ffi::lua_resetthread(thread_state); - reset_ok = true; - } - - if reset_ok { + if extra.thread_pool.len() < extra.thread_pool.capacity() { extra.thread_pool.push(thread.0.index); thread.0.drop = false; // Prevent thread from being garbage collected } @@ -1244,7 +1209,7 @@ impl RawLua { let rawlua = (*extra).raw_lua(); let func = &*(*upvalue).data; - let fut = func(rawlua, nargs); + let fut = Some(func(rawlua, nargs)); let extra = XRc::clone(&(*upvalue).extra); let protect = !rawlua.unlikely_memory_error(); push_internal_userdata(state, AsyncPollUpvalue { data: fut, extra }, protect)?; @@ -1262,20 +1227,27 @@ impl RawLua { unsafe extern "C-unwind" fn poll_future(state: *mut ffi::lua_State) -> c_int { let upvalue = get_userdata::(state, ffi::lua_upvalueindex(1)); - callback_error_ext(state, (*upvalue).extra.get(), true, |extra, _| { + callback_error_ext(state, (*upvalue).extra.get(), true, |extra, nargs| { // Lua ensures that `LUA_MINSTACK` stack spaces are available (after pushing arguments) // The lock must be already held as the future is polled let rawlua = (*extra).raw_lua(); + if nargs == 1 && ffi::lua_tolightuserdata(state, -1) == Lua::poll_terminate().0 { + // Destroy the future and terminate the Lua thread + (*upvalue).data.take(); + ffi::lua_pushinteger(state, 0); + return Ok(1); + } + let fut = &mut (*upvalue).data; let mut ctx = Context::from_waker(rawlua.waker()); - match fut.as_mut().poll(&mut ctx) { - Poll::Pending => { + match fut.as_mut().map(|fut| fut.as_mut().poll(&mut ctx)) { + Some(Poll::Pending) => { ffi::lua_pushnil(state); ffi::lua_pushlightuserdata(state, Lua::poll_pending().0); Ok(2) } - Poll::Ready(nresults) => { + Some(Poll::Ready(nresults)) => { match nresults? { nresults if nresults < 3 => { // Fast path for up to 2 results without creating a table @@ -1293,6 +1265,7 @@ impl RawLua { } } } + None => Err(Error::CallbackDestructed), } }) } @@ -1338,8 +1311,8 @@ impl RawLua { lua.load( r#" local poll = get_poll(...) + local nres, res, res2 = poll() while true do - local nres, res, res2 = poll() if nres ~= nil then if nres == 0 then return @@ -1351,7 +1324,10 @@ impl RawLua { return unpack(res, nres) end end - yield(res) -- `res` is a "pending" value + -- `res` is a "pending" value + -- `yield` can return a signal to drop the future that we should propagate + -- to the poller + nres, res, res2 = poll(yield(res)) end "#, ) diff --git a/src/thread.rs b/src/thread.rs index 39394ce..da7962b 100644 --- a/src/thread.rs +++ b/src/thread.rs @@ -305,29 +305,10 @@ impl Thread { pub fn reset(&self, func: Function) -> Result<()> { let lua = self.0.lua.lock(); let thread_state = self.state(); - match self.status_inner(&lua) { - ThreadStatusInner::Running => return Err(Error::runtime("cannot reset a running thread")), - // Any Lua can reuse new or finished thread - ThreadStatusInner::New(_) => unsafe { ffi::lua_settop(thread_state, 0) }, - ThreadStatusInner::Finished => {} - #[cfg(not(any(feature = "lua54", feature = "luau")))] - _ => return Err(Error::runtime("cannot reset non-finished thread")), - #[cfg(any(feature = "lua54", feature = "luau"))] - _ => unsafe { - #[cfg(all(feature = "lua54", not(feature = "vendored")))] - let status = ffi::lua_resetthread(thread_state); - #[cfg(all(feature = "lua54", feature = "vendored"))] - let status = ffi::lua_closethread(thread_state, lua.state()); - #[cfg(feature = "lua54")] - if status != ffi::LUA_OK { - return Err(pop_error(thread_state, status)); - } - #[cfg(feature = "luau")] - ffi::lua_resetthread(thread_state); - }, - } - unsafe { + let status = self.status_inner(&lua); + self.reset_inner(status)?; + // Push function to the top of the thread stack ffi::lua_xpush(lua.ref_thread(), thread_state, func.0.index); @@ -342,6 +323,42 @@ impl Thread { } } + unsafe fn reset_inner(&self, status: ThreadStatusInner) -> Result<()> { + match status { + ThreadStatusInner::New(_) => { + // The thread is new, so we can just set the top to 0 + ffi::lua_settop(self.state(), 0); + Ok(()) + } + ThreadStatusInner::Running => Err(Error::runtime("cannot reset a running thread")), + ThreadStatusInner::Finished => Ok(()), + #[cfg(not(any(feature = "lua54", feature = "luau")))] + ThreadStatusInner::Yielded(_) | ThreadStatusInner::Error => { + Err(Error::runtime("cannot reset non-finished thread")) + } + #[cfg(any(feature = "lua54", feature = "luau"))] + ThreadStatusInner::Yielded(_) | ThreadStatusInner::Error => { + let thread_state = self.state(); + + #[cfg(all(feature = "lua54", not(feature = "vendored")))] + let status = ffi::lua_resetthread(thread_state); + #[cfg(all(feature = "lua54", feature = "vendored"))] + let status = { + let lua = self.0.lua.lock(); + ffi::lua_closethread(thread_state, lua.state()) + }; + #[cfg(feature = "lua54")] + if status != ffi::LUA_OK { + return Err(pop_error(thread_state, status)); + } + #[cfg(feature = "luau")] + ffi::lua_resetthread(thread_state); + + Ok(()) + } + } + } + /// Converts [`Thread`] to an [`AsyncThread`] which implements [`Future`] and [`Stream`] traits. /// /// Only resumable threads can be converted to [`AsyncThread`]. @@ -505,8 +522,21 @@ impl Drop for AsyncThread { fn drop(&mut self) { if self.recycle { if let Some(lua) = self.thread.0.lua.try_lock() { - // For Lua 5.4 this also closes all pending to-be-closed variables - unsafe { lua.recycle_thread(&mut self.thread) }; + unsafe { + let mut status = self.thread.status_inner(&lua); + if matches!(status, ThreadStatusInner::Yielded(0)) { + // The thread is dropped while yielded, resume it with the "terminate" signal + ffi::lua_pushlightuserdata(self.thread.1, crate::Lua::poll_terminate().0); + if let Ok((new_status, _)) = self.thread.resume_inner(&lua, 1) { + status = new_status; + } + } + + // For Lua 5.4 this also closes all pending to-be-closed variables + if self.thread.reset_inner(status).is_ok() { + lua.recycle_thread(&mut self.thread); + } + } } } } diff --git a/src/types.rs b/src/types.rs index 2f63906..2589ea6 100644 --- a/src/types.rs +++ b/src/types.rs @@ -61,7 +61,7 @@ pub(crate) type AsyncCallback = pub(crate) type AsyncCallbackUpvalue = Upvalue; #[cfg(feature = "async")] -pub(crate) type AsyncPollUpvalue = Upvalue>>; +pub(crate) type AsyncPollUpvalue = Upvalue>>>; /// Type to set next Lua VM action after executing interrupt or hook function. pub enum VmState { diff --git a/tests/async.rs b/tests/async.rs index 6ae3b3d..752abb5 100644 --- a/tests/async.rs +++ b/tests/async.rs @@ -9,7 +9,7 @@ use tokio::sync::Mutex; use mlua::{ Error, Function, Lua, LuaOptions, MultiValue, ObjectLike, Result, StdLib, Table, UserData, - UserDataMethods, Value, + UserDataMethods, UserDataRef, Value, }; #[cfg(not(target_arch = "wasm32"))] @@ -547,6 +547,7 @@ async fn test_async_thread_error() -> Result<()> { #[tokio::test] async fn test_async_terminate() -> Result<()> { + // Future captures `Lua` instance and dropped all together let mutex = Arc::new(Mutex::new(0u32)); { let lua = Lua::new(); @@ -565,6 +566,17 @@ async fn test_async_terminate() -> Result<()> { } assert!(mutex.try_lock().is_ok()); + // Future is dropped, but `Lua` instance is still alive + let lua = Lua::new(); + let func = lua.create_async_function(move |_, mutex: UserDataRef>>| async move { + let _guard = mutex.lock().await; + sleep_ms(100).await; + Ok(()) + })?; + let mutex2 = lua.create_any_userdata(mutex.clone())?; + let _ = tokio::time::timeout(Duration::from_millis(30), func.call_async::<()>(mutex2)).await; + assert!(mutex.try_lock().is_ok()); + Ok(()) }