Fix coroutine stack handling after hook yields

Preserve the coroutine stack when a debug hook yields, so Lua stack
is not truncated by `StackGuard`.

Fixes #723

Author: khvzak

src/state.rs

@@ -782,7 +782,7 @@ impl Lua {
     pub fn remove_hook(&self) {
         let lua = self.lock();
         unsafe {
-            ffi::lua_sethook(lua.state(), None, 0, 0);
+            lua.remove_thread_hook(lua.state());
         }
     }
 

src/state/extra.rs

@@ -81,6 +81,8 @@ pub(crate) struct ExtraData {
     pub(super) hook_callback: Option<crate::types::HookCallback>,
     #[cfg(not(feature = "luau"))]
     pub(super) hook_triggers: crate::debug::HookTriggers,
+    #[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+    pub(super) hook_removed_while_yielded: bool,
     #[cfg(any(feature = "lua55", feature = "lua54"))]
     pub(super) warn_callback: Option<crate::types::WarnCallback>,
     #[cfg(feature = "luau")]
@@ -187,6 +189,8 @@ impl ExtraData {
             hook_callback: None,
             #[cfg(not(feature = "luau"))]
             hook_triggers: Default::default(),
+            #[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+            hook_removed_while_yielded: false,
             #[cfg(any(feature = "lua55", feature = "lua54"))]
             warn_callback: None,
             #[cfg(feature = "luau")]

src/state/raw.rs

@@ -427,6 +427,10 @@ impl RawLua {
                     if event == ffi::LUA_HOOKCOUNT || event == ffi::LUA_HOOKLINE {
                         #[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
                         if ffi::lua_isyieldable(state) != 0 {
+                            if ffi::lua_gethook(state).is_none() {
+                                let extra = ExtraData::get(state);
+                                (*extra).hook_removed_while_yielded = true;
+                            }
                             ffi::lua_yield(state, 0);
                         }
                         #[cfg(any(feature = "lua52", feature = "lua51", feature = "luajit"))]
@@ -518,6 +522,41 @@ impl RawLua {
         Ok(())
     }
 
+    #[cfg(not(feature = "luau"))]
+    #[inline]
+    pub(crate) unsafe fn remove_thread_hook(&self, thread_state: *mut ffi::lua_State) {
+        #[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+        if ffi::lua_status(thread_state) == ffi::LUA_YIELD && Self::has_hook_yielded_frame(thread_state) {
+            (*self.extra.get()).hook_removed_while_yielded = true;
+        }
+        ffi::lua_sethook(thread_state, None, 0, 0);
+    }
+
+    #[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+    unsafe fn has_hook_yielded_frame(thread_state: *mut ffi::lua_State) -> bool {
+        let mut ar = mem::zeroed::<ffi::lua_Debug>();
+        ffi::lua_getstack(thread_state, 0, &mut ar) != 0
+            && ffi::lua_getinfo(thread_state, cstr!("S"), &mut ar) != 0
+            && !ar.what.is_null()
+            && CStr::from_ptr(ar.what).to_bytes() != b"C"
+    }
+
+    pub(crate) unsafe fn is_hook_yielded(&self, thread_state: *mut ffi::lua_State) -> bool {
+        #[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+        {
+            if ffi::lua_gethook(thread_state).is_none() && !(*self.extra.get()).hook_removed_while_yielded {
+                return false;
+            }
+            Self::has_hook_yielded_frame(thread_state)
+        }
+
+        #[cfg(not(any(feature = "lua55", feature = "lua54", feature = "lua53")))]
+        {
+            let _ = thread_state;
+            false
+        }
+    }
+
     /// See [`Lua::create_string`]
     pub(crate) unsafe fn create_string(&self, s: &[u8]) -> Result<LuaString> {
         let state = self.state();

src/thread.rs

@@ -185,12 +185,6 @@ enum ThreadStatusInner {
 }
 
 impl ThreadStatusInner {
-    #[cfg(feature = "async")]
-    #[inline(always)]
-    fn is_resumable(self) -> bool {
-        matches!(self, ThreadStatusInner::New(_) | ThreadStatusInner::Yielded(_))
-    }
-
     #[inline(always)]
     fn is_yielded(self) -> bool {
         matches!(self, ThreadStatusInner::Yielded(_))
@@ -329,7 +323,7 @@ impl Thread {
     {
         let lua = self.0.lua.lock();
         check_thread_reentrancy(self.state(), &lua)?;
-        let mut pushed_nargs = self.resumable_nargs(&lua)?;
+        let (mut pushed_nargs, mut hook_yielded) = self.resumable_state(&lua)?;
 
         let state = lua.state();
         let thread_state = self.state();
@@ -341,18 +335,24 @@ impl Thread {
             if exec_thread_event(&lua, on_resume, thread_state, || {
                 ThreadEvent::Resume(self.clone())
             })? {
-                pushed_nargs = self.resumable_nargs(&lua)?;
+                (pushed_nargs, hook_yielded) = self.resumable_state(&lua)?;
             }
 
-            let nargs = args.push_into_stack_multi(&lua)?;
-            if nargs > 0 {
-                check_stack(thread_state, nargs)?;
-                ffi::lua_xmove(state, thread_state, nargs);
-                pushed_nargs += nargs;
+            if !hook_yielded {
+                let nargs = args.push_into_stack_multi(&lua)?;
+                if nargs > 0 {
+                    check_stack(thread_state, nargs)?;
+                    ffi::lua_xmove(state, thread_state, nargs);
+                    pushed_nargs += nargs;
+                }
             }
 
-            let _thread_sg = StackGuard::with_top(thread_state, 0);
+            let mut thread_sg = StackGuard::with_top(thread_state, 0);
             let (status, nresults) = self.resume_inner(&lua, pushed_nargs)?;
+            if status.is_yielded() && self.is_hook_yielded(&lua) {
+                debug_assert_eq!(nresults, 0);
+                thread_sg.keep(ffi::lua_gettop(thread_state));
+            }
 
             check_stack(state, nresults + 1)?;
             ffi::lua_xmove(thread_state, state, nresults);
@@ -478,15 +478,24 @@ impl Thread {
         }
     }
 
-    /// Returns the number of pending arguments on the thread stack if the thread is resumable.
+    /// Returns the pending argument count and whether the thread was interrupted by a hook.
     #[inline]
-    fn resumable_nargs(&self, lua: &RawLua) -> Result<c_int> {
+    fn resumable_state(&self, lua: &RawLua) -> Result<(c_int, bool)> {
         match self.status_inner(lua) {
-            ThreadStatusInner::New(nargs) | ThreadStatusInner::Yielded(nargs) => Ok(nargs),
+            ThreadStatusInner::New(nargs) => Ok((nargs, false)),
+            ThreadStatusInner::Yielded(nargs) => {
+                let hook_yielded = self.is_hook_yielded(lua);
+                Ok((if hook_yielded { 0 } else { nargs }, hook_yielded))
+            }
             _ => Err(Error::CoroutineUnresumable),
         }
     }
 
+    /// Distinguishes a hook interruption from a normal yield.
+    fn is_hook_yielded(&self, lua: &RawLua) -> bool {
+        unsafe { lua.is_hook_yielded(self.state()) }
+    }
+
     /// Returns `true` if this thread is resumable (meaning it can be resumed by calling
     /// [`Thread::resume`]).
     #[inline(always)]
@@ -548,9 +557,9 @@ impl Thread {
     #[cfg(not(feature = "luau"))]
     #[cfg_attr(docsrs, doc(cfg(not(feature = "luau"))))]
     pub fn remove_hook(&self) {
-        let _lua = self.0.lua.lock();
+        let lua = self.0.lua.lock();
         unsafe {
-            ffi::lua_sethook(self.state(), None, 0, 0);
+            lua.remove_thread_hook(self.state());
         }
     }
 
@@ -681,19 +690,19 @@ impl Thread {
     {
         let lua = self.0.lua.lock();
         check_thread_reentrancy(self.state(), &lua)?;
-        if !self.status_inner(&lua).is_resumable() {
-            return Err(Error::CoroutineUnresumable);
-        }
+        let (_, hook_yielded) = self.resumable_state(&lua)?;
 
         let state = lua.state();
         let thread_state = self.state();
         unsafe {
             let _sg = StackGuard::new(state);
 
-            let nargs = args.push_into_stack_multi(&lua)?;
-            if nargs > 0 {
-                check_stack(thread_state, nargs)?;
-                ffi::lua_xmove(state, thread_state, nargs);
+            if !hook_yielded {
+                let nargs = args.push_into_stack_multi(&lua)?;
+                if nargs > 0 {
+                    check_stack(thread_state, nargs)?;
+                    ffi::lua_xmove(state, thread_state, nargs);
+                }
             }
 
             Ok(AsyncThread {
@@ -794,7 +803,7 @@ impl<R> Drop for AsyncThread<R> {
         {
             unsafe {
                 let mut status = self.thread.status_inner(&lua);
-                if matches!(status, ThreadStatusInner::Yielded(0)) {
+                if matches!(status, ThreadStatusInner::Yielded(0)) && !self.thread.is_hook_yielded(&lua) {
                     // 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) {
@@ -820,32 +829,37 @@ impl<R: FromLuaMulti> Stream for AsyncThread<R> {
     fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
         let lua = self.thread.0.lua.lock();
         check_thread_reentrancy(self.thread.state(), &lua)?;
-        let mut nargs = match self.thread.resumable_nargs(&lua) {
-            Ok(nargs) => nargs,
+        let mut nargs = match self.thread.resumable_state(&lua) {
+            Ok((nargs, _)) => nargs,
             Err(_) => return Poll::Ready(None),
         };
 
         let state = lua.state();
         let thread_state = self.thread.state();
         unsafe {
             let _sg = StackGuard::new(state);
-            let _thread_sg = StackGuard::with_top(thread_state, 0);
+            let mut thread_sg = StackGuard::with_top(thread_state, 0);
             let _wg = WakerGuard::new(&lua, cx.waker());
 
             // If the resume callback runs, it may touch this thread, so re-read the argument count
             let on_resume = lua.thread_event_triggers().on_resume;
             if exec_thread_event(&lua, on_resume, thread_state, || {
                 ThreadEvent::Resume(self.thread.clone())
             })? {
-                nargs = match self.thread.resumable_nargs(&lua) {
-                    Ok(nargs) => nargs,
+                nargs = match self.thread.resumable_state(&lua) {
+                    Ok((nargs, _)) => nargs,
                     Err(_) => return Poll::Ready(None),
                 };
             }
 
             let (status, nresults) = (self.thread).resume_inner(&lua, nargs)?;
+            let hook_yielded = status.is_yielded() && self.thread.is_hook_yielded(&lua);
+            if hook_yielded {
+                debug_assert_eq!(nresults, 0);
+                thread_sg.keep(ffi::lua_gettop(thread_state));
+            }
 
-            if status.is_yielded() && nresults == 1 && is_poll_pending(thread_state) {
+            if status.is_yielded() && !hook_yielded && nresults == 1 && is_poll_pending(thread_state) {
                 // Exec thread yield callback
                 let on_yield = lua.thread_event_triggers().on_yield;
                 exec_thread_event(&lua, on_yield, thread_state, || {
@@ -878,27 +892,32 @@ impl<R: FromLuaMulti> Future for AsyncThread<R> {
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let lua = self.thread.0.lua.lock();
         check_thread_reentrancy(self.thread.state(), &lua)?;
-        let mut nargs = self.thread.resumable_nargs(&lua)?;
+        let (mut nargs, _) = self.thread.resumable_state(&lua)?;
 
         let state = lua.state();
         let thread_state = self.thread.state();
         unsafe {
             let _sg = StackGuard::new(state);
-            let _thread_sg = StackGuard::with_top(thread_state, 0);
+            let mut thread_sg = StackGuard::with_top(thread_state, 0);
             let _wg = WakerGuard::new(&lua, cx.waker());
 
             // If the resume callback runs, it may touch this thread, so re-read the argument count
             let on_resume = lua.thread_event_triggers().on_resume;
             if exec_thread_event(&lua, on_resume, thread_state, || {
                 ThreadEvent::Resume(self.thread.clone())
             })? {
-                nargs = self.thread.resumable_nargs(&lua)?;
+                (nargs, _) = self.thread.resumable_state(&lua)?;
             }
 
             let (status, nresults) = self.thread.resume_inner(&lua, nargs)?;
+            let hook_yielded = status.is_yielded() && self.thread.is_hook_yielded(&lua);
+            if hook_yielded {
+                debug_assert_eq!(nresults, 0);
+                thread_sg.keep(ffi::lua_gettop(thread_state));
+            }
 
             if status.is_yielded() {
-                let pending = nresults == 1 && is_poll_pending(thread_state);
+                let pending = !hook_yielded && nresults == 1 && is_poll_pending(thread_state);
 
                 // Exec thread yield callback
                 let on_yield = lua.thread_event_triggers().on_yield;

tests/async.rs

@@ -687,6 +687,49 @@ async fn test_async_hook() -> Result<()> {
     Ok(())
 }
 
+#[tokio::test]
+#[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+async fn test_async_hook_yield_preserves_stack() -> Result<()> {
+    use std::future::{Future, poll_fn};
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::task::Poll;
+
+    let lua = Lua::new();
+
+    let thread = lua.create_thread(
+        lua.load(
+            r#"
+            local x = 40
+            local y = 2
+            return x + y
+        "#,
+        )
+        .into_function()?,
+    )?;
+
+    let yielded = Arc::new(AtomicBool::new(false));
+    let yielded2 = yielded.clone();
+    thread.set_hook(mlua::HookTriggers::EVERY_LINE, move |lua, debug| {
+        if debug.current_line() == Some(4) && !yielded2.swap(true, Ordering::Relaxed) {
+            lua.remove_hook();
+            return Ok(mlua::VmState::Yield);
+        }
+        Ok(mlua::VmState::Continue)
+    })?;
+
+    let mut thread = Box::pin(thread.into_async::<i32>(())?);
+    poll_fn(|cx| {
+        assert!(thread.as_mut().poll(cx).is_pending());
+        Poll::Ready(())
+    })
+    .await;
+    assert!(yielded.load(Ordering::Relaxed));
+    lua.gc_collect()?;
+    assert_eq!(thread.await?, 42);
+
+    Ok(())
+}
+
 #[test]
 fn test_async_yield_with() -> Result<()> {
     let lua = Lua::new();

tests/hooks.rs

@@ -294,6 +294,40 @@ fn test_hook_yield() -> Result<()> {
     Ok(())
 }
 
+#[test]
+#[cfg(any(feature = "lua55", feature = "lua54", feature = "lua53"))]
+fn test_hook_yield_preserves_stack() -> Result<()> {
+    let lua = Lua::new();
+
+    let func = lua
+        .load(
+            r#"
+            local x = { value = 40 }
+            local y = 2
+            return x.value + y
+        "#,
+        )
+        .into_function()?;
+    let co = lua.create_thread(func)?;
+
+    let yielded = Arc::new(std::sync::atomic::AtomicBool::new(false));
+    let yielded2 = yielded.clone();
+    co.set_hook(HookTriggers::EVERY_LINE, move |_lua, debug| {
+        if debug.current_line() == Some(4) && !yielded2.swap(true, Ordering::Relaxed) {
+            return Ok(VmState::Yield);
+        }
+        Ok(VmState::Continue)
+    })?;
+
+    co.resume::<()>(())?;
+    assert!(yielded.load(Ordering::Relaxed));
+    co.remove_hook();
+    lua.gc_collect()?;
+    assert_eq!(co.resume::<i32>(())?, 42);
+
+    Ok(())
+}
+
 #[test]
 fn test_global_hook() -> Result<()> {
     let lua = Lua::new();