Enable Thread::reset() for all Lua versions

khvzak · khvzak · commit bcb9a4d220a5 · 2025-05-04T12:07:55.000+01:00
diff --git a/src/state.rs b/src/state.rs
@@ -98,9 +98,6 @@ pub struct LuaOptions {
 
     /// Max size of thread (coroutine) object pool used to execute asynchronous functions.
     ///
-    /// It works on Lua 5.4 and Luau, where [`lua_resetthread`] function
-    /// is available and allows to reuse old coroutines after resetting their state.
-    ///
     /// Default: **0** (disabled)
     ///
     /// [`lua_resetthread`]: https://www.lua.org/manual/5.4/manual.html#lua_resetthread
diff --git a/src/state/raw.rs b/src/state/raw.rs
@@ -505,7 +505,6 @@ impl RawLua {
     /// Wraps a Lua function into a new or recycled thread (coroutine).
     #[cfg(feature = "async")]
     pub(crate) unsafe fn create_recycled_thread(&self, func: &Function) -> Result<Thread> {
-        #[cfg(any(feature = "lua54", feature = "luau"))]
         if let Some(index) = (*self.extra.get()).thread_pool.pop() {
             let thread_state = ffi::lua_tothread(self.ref_thread(), index);
             ffi::lua_xpush(self.ref_thread(), thread_state, func.0.index);
@@ -525,27 +524,47 @@ impl RawLua {
 
     /// Resets thread (coroutine) and returns it to the pool for later use.
     #[cfg(feature = "async")]
-    #[cfg(any(feature = "lua54", feature = "luau"))]
-    pub(crate) unsafe fn recycle_thread(&self, thread: &mut Thread) -> bool {
+    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() {
-            let thread_state = ffi::lua_tothread(extra.ref_thread, thread.0.index);
-            #[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, self.state());
+        if extra.thread_pool.len() == extra.thread_pool.capacity() {
             #[cfg(feature = "lua54")]
-            if status != ffi::LUA_OK {
-                // Error object is on top, drop it
+            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);
             }
-            #[cfg(feature = "luau")]
+            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 {
             extra.thread_pool.push(thread.0.index);
             thread.0.drop = false; // Prevent thread from being garbage collected
-            return true;
         }
-        false
     }
 
     /// Pushes a value that implements `IntoLua` onto the Lua stack.
diff --git a/src/thread.rs b/src/thread.rs
@@ -2,8 +2,7 @@ use std::fmt;
 use std::os::raw::{c_int, c_void};
 
 use crate::error::{Error, Result};
-#[allow(unused)]
-use crate::state::Lua;
+use crate::function::Function;
 use crate::state::RawLua;
 use crate::traits::{FromLuaMulti, IntoLuaMulti};
 use crate::types::{LuaType, ValueRef};
@@ -42,6 +41,26 @@ pub enum ThreadStatus {
     Error,
 }
 
+/// Internal representation of a Lua thread status.
+///
+/// The number in `New` and `Yielded` variants is the number of arguments pushed
+/// to the thread stack.
+#[derive(Clone, Copy)]
+enum ThreadStatusInner {
+    New,
+    Running,
+    Yielded,
+    Finished,
+    Error,
+}
+
+impl ThreadStatusInner {
+    #[inline(always)]
+    fn is_resumable(self) -> bool {
+        matches!(self, ThreadStatusInner::New | ThreadStatusInner::Yielded)
+    }
+}
+
 /// Handle to an internal Lua thread (coroutine).
 #[derive(Clone)]
 pub struct Thread(pub(crate) ValueRef, pub(crate) *mut ffi::lua_State);
@@ -122,7 +141,7 @@ impl Thread {
         R: FromLuaMulti,
     {
         let lua = self.0.lua.lock();
-        if self.status_inner(&lua) != ThreadStatus::Resumable {
+        if !self.status_inner(&lua).is_resumable() {
             return Err(Error::CoroutineUnresumable);
         }
 
@@ -170,23 +189,27 @@ impl Thread {
 
     /// Gets the status of the thread.
     pub fn status(&self) -> ThreadStatus {
-        self.status_inner(&self.0.lua.lock())
+        match self.status_inner(&self.0.lua.lock()) {
+            ThreadStatusInner::New | ThreadStatusInner::Yielded => ThreadStatus::Resumable,
+            ThreadStatusInner::Running => ThreadStatus::Running,
+            ThreadStatusInner::Finished => ThreadStatus::Finished,
+            ThreadStatusInner::Error => ThreadStatus::Error,
+        }
     }
 
     /// Gets the status of the thread (internal implementation).
-    pub(crate) fn status_inner(&self, lua: &RawLua) -> ThreadStatus {
+    fn status_inner(&self, lua: &RawLua) -> ThreadStatusInner {
         let thread_state = self.state();
         if thread_state == lua.state() {
             // The thread is currently running
-            return ThreadStatus::Running;
+            return ThreadStatusInner::Running;
         }
         let status = unsafe { ffi::lua_status(thread_state) };
-        if status != ffi::LUA_OK && status != ffi::LUA_YIELD {
-            ThreadStatus::Error
-        } else if status == ffi::LUA_YIELD || unsafe { ffi::lua_gettop(thread_state) > 0 } {
-            ThreadStatus::Resumable
-        } else {
-            ThreadStatus::Finished
+        match status {
+            ffi::LUA_YIELD => ThreadStatusInner::Yielded,
+            ffi::LUA_OK if unsafe { ffi::lua_gettop(thread_state) } > 0 => ThreadStatusInner::New,
+            ffi::LUA_OK => ThreadStatusInner::Finished,
+            _ => ThreadStatusInner::Error,
         }
     }
 
@@ -198,7 +221,7 @@ impl Thread {
     #[cfg_attr(docsrs, doc(cfg(not(feature = "luau"))))]
     pub fn set_hook<F>(&self, triggers: HookTriggers, callback: F)
     where
-        F: Fn(&Lua, Debug) -> Result<crate::VmState> + MaybeSend + 'static,
+        F: Fn(&crate::Lua, Debug) -> Result<crate::VmState> + MaybeSend + 'static,
     {
         let lua = self.0.lua.lock();
         unsafe {
@@ -215,32 +238,37 @@ impl Thread {
     /// In Luau: resets to the initial state of a newly created Lua thread.
     /// Lua threads in arbitrary states (like yielded or errored) can be reset properly.
     ///
-    /// Sets a Lua function for the thread afterwards.
+    /// Other Lua versions can reset only new or finished threads.
     ///
-    /// Requires `feature = "lua54"` OR `feature = "luau"`.
+    /// Sets a Lua function for the thread afterwards.
     ///
     /// [Lua 5.4]: https://www.lua.org/manual/5.4/manual.html#lua_closethread
-    #[cfg(any(feature = "lua54", feature = "luau"))]
-    #[cfg_attr(docsrs, doc(cfg(any(feature = "lua54", feature = "luau"))))]
-    pub fn reset(&self, func: crate::function::Function) -> Result<()> {
+    pub fn reset(&self, func: Function) -> Result<()> {
         let lua = self.0.lua.lock();
-        if self.status_inner(&lua) == ThreadStatus::Running {
-            return Err(Error::runtime("cannot reset a running thread"));
+        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);
+            },
         }
 
-        let thread_state = self.state();
         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);
-
             // Push function to the top of the thread stack
             ffi::lua_xpush(lua.ref_thread(), thread_state, func.0.index);
 
@@ -393,30 +421,19 @@ impl LuaType for Thread {
 
 #[cfg(feature = "async")]
 impl<A, R> AsyncThread<A, R> {
-    #[inline]
+    #[inline(always)]
     pub(crate) fn set_recyclable(&mut self, recyclable: bool) {
         self.recycle = recyclable;
     }
 }
 
 #[cfg(feature = "async")]
-#[cfg(any(feature = "lua54", feature = "luau"))]
 impl<A, R> Drop for AsyncThread<A, R> {
     fn drop(&mut self) {
         if self.recycle {
             if let Some(lua) = self.thread.0.lua.try_lock() {
-                unsafe {
-                    // For Lua 5.4 this also closes all pending to-be-closed variables
-                    if !lua.recycle_thread(&mut self.thread) {
-                        #[cfg(feature = "lua54")]
-                        if self.thread.status_inner(&lua) == ThreadStatus::Error {
-                            #[cfg(not(feature = "vendored"))]
-                            ffi::lua_resetthread(self.thread.state());
-                            #[cfg(feature = "vendored")]
-                            ffi::lua_closethread(self.thread.state(), lua.state());
-                        }
-                    }
-                }
+                // For Lua 5.4 this also closes all pending to-be-closed variables
+                unsafe { lua.recycle_thread(&mut self.thread) };
             }
         }
     }
@@ -428,7 +445,7 @@ impl<A: IntoLuaMulti, R: FromLuaMulti> Stream for AsyncThread<A, R> {
 
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let lua = self.thread.0.lua.lock();
-        if self.thread.status_inner(&lua) != ThreadStatus::Resumable {
+        if !self.thread.status_inner(&lua).is_resumable() {
             return Poll::Ready(None);
         }
 
@@ -466,7 +483,7 @@ impl<A: IntoLuaMulti, R: FromLuaMulti> Future for AsyncThread<A, R> {
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let lua = self.thread.0.lua.lock();
-        if self.thread.status_inner(&lua) != ThreadStatus::Resumable {
+        if !self.thread.status_inner(&lua).is_resumable() {
             return Poll::Ready(Err(Error::CoroutineUnresumable));
         }
 
@@ -506,7 +523,7 @@ impl<A: IntoLuaMulti, R: FromLuaMulti> Future for AsyncThread<A, R> {
 #[cfg(feature = "async")]
 #[inline(always)]
 unsafe fn is_poll_pending(state: *mut ffi::lua_State) -> bool {
-    ffi::lua_tolightuserdata(state, -1) == Lua::poll_pending().0
+    ffi::lua_tolightuserdata(state, -1) == crate::Lua::poll_pending().0
 }
 
 #[cfg(feature = "async")]
diff --git a/tests/thread.rs b/tests/thread.rs
@@ -107,7 +107,6 @@ fn test_thread() -> Result<()> {
 }
 
 #[test]
-#[cfg(any(feature = "lua54", feature = "luau"))]
 fn test_thread_reset() -> Result<()> {
     use mlua::{AnyUserData, UserData};
     use std::sync::Arc;
@@ -120,7 +119,8 @@ fn test_thread_reset() -> Result<()> {
     let arc = Arc::new(());
 
     let func: Function = lua.load(r#"function(ud) coroutine.yield(ud) end"#).eval()?;
-    let thread = lua.create_thread(func.clone())?;
+    let thread = lua.create_thread(lua.load("return 0").into_function()?)?; // Dummy function first
+    assert!(thread.reset(func.clone()).is_ok());
 
     for _ in 0..2 {
         assert_eq!(thread.status(), ThreadStatus::Resumable);
@@ -145,11 +145,7 @@ fn test_thread_reset() -> Result<()> {
         assert!(thread.reset(func.clone()).is_err());
         // Reset behavior has changed in Lua v5.4.4
         // It's became possible to force reset thread by popping error object
-        assert!(matches!(
-            thread.status(),
-            ThreadStatus::Finished | ThreadStatus::Error
-        ));
-        // Would pass in 5.4.4
+        assert!(matches!(thread.status(), ThreadStatus::Finished));
         assert!(thread.reset(func.clone()).is_ok());
         assert_eq!(thread.status(), ThreadStatus::Resumable);
     }
