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.

Author: khvzak

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

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::<AsyncPollUpvalue>(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
             "#,
         )

src/thread.rs

@@ -305,40 +305,57 @@ 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 => {}
+        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);
+
+            #[cfg(feature = "luau")]
+            {
+                // Inherit `LUA_GLOBALSINDEX` from the main thread
+                ffi::lua_xpush(lua.main_state(), thread_state, ffi::LUA_GLOBALSINDEX);
+                ffi::lua_replace(thread_state, ffi::LUA_GLOBALSINDEX);
+            }
+
+            Ok(())
+        }
+    }
+
+    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")))]
-            _ => return Err(Error::runtime("cannot reset non-finished thread")),
+            ThreadStatusInner::Yielded(_) | ThreadStatusInner::Error => {
+                Err(Error::runtime("cannot reset non-finished thread"))
+            }
             #[cfg(any(feature = "lua54", feature = "luau"))]
-            _ => unsafe {
+            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 = ffi::lua_closethread(thread_state, lua.state());
+                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);
-            },
-        }
-
-        unsafe {
-            // Push function to the top of the thread stack
-            ffi::lua_xpush(lua.ref_thread(), thread_state, func.0.index);
 
-            #[cfg(feature = "luau")]
-            {
-                // Inherit `LUA_GLOBALSINDEX` from the main thread
-                ffi::lua_xpush(lua.main_state(), thread_state, ffi::LUA_GLOBALSINDEX);
-                ffi::lua_replace(thread_state, ffi::LUA_GLOBALSINDEX);
+                Ok(())
             }
-
-            Ok(())
         }
     }
 
@@ -505,8 +522,21 @@ impl<R> Drop for AsyncThread<R> {
     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);
+                    }
+                }
             }
         }
     }

src/types.rs

@@ -61,7 +61,7 @@ pub(crate) type AsyncCallback =
 pub(crate) type AsyncCallbackUpvalue = Upvalue<AsyncCallback>;
 
 #[cfg(feature = "async")]
-pub(crate) type AsyncPollUpvalue = Upvalue<BoxFuture<'static, Result<c_int>>>;
+pub(crate) type AsyncPollUpvalue = Upvalue<Option<BoxFuture<'static, Result<c_int>>>>;
 
 /// Type to set next Lua VM action after executing interrupt or hook function.
 pub enum VmState {

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<Arc<Mutex<u32>>>| 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(())
 }