Map implicit async threads to their root owner

`Lua::current_thread()` now returns a stable handle instead of a temporary coroutine.
Closes #706

Author: khvzak

src/function.rs

@@ -252,6 +252,7 @@ impl Function {
             lua.create_recycled_thread(self).and_then(|th| {
                 let mut th = th.into_async(args)?;
                 th.set_recyclable(true);
+                lua.update_thread_ownership(th.thread(), Some(lua.state()));
                 Ok(th)
             })
         })

src/state.rs

@@ -1764,7 +1764,18 @@ impl Lua {
     pub fn current_thread(&self) -> Thread {
         let lua = self.lock();
         let state = lua.state();
+        let extra = lua.extra.get();
         unsafe {
+            // If this thread is implicit (created by `call_async`), return the root user-owned thread
+            // instead.
+            #[cfg(feature = "async")]
+            if let Some(&owner) = (*extra).thread_ownership_map.get(&state) {
+                assert_stack(owner, 1);
+                ffi::lua_pushthread(owner);
+                ffi::lua_xmove(owner, lua.ref_thread(), 1);
+                return Thread(lua.pop_ref_thread(), owner);
+            }
+
             let _sg = StackGuard::new(state);
             assert_stack(state, 1);
             ffi::lua_pushthread(state);

src/state/extra.rs

@@ -66,6 +66,9 @@ pub(crate) struct ExtraData {
     // Pool of `Thread`s (coroutines) for async execution
     #[cfg(feature = "async")]
     pub(super) thread_pool: Vec<crate::types::ValueRefIndex>,
+    // Map for implicit threads to root user-owned Thread
+    #[cfg(feature = "async")]
+    pub(super) thread_ownership_map: FxHashMap<*mut ffi::lua_State, *mut ffi::lua_State>,
 
     // Address of `WrappedFailure` metatable
     pub(super) wrapped_failure_mt_ptr: *const c_void,
@@ -174,6 +177,8 @@ impl ExtraData {
             wrapped_failure_top: 0,
             #[cfg(feature = "async")]
             thread_pool: Vec::new(),
+            #[cfg(feature = "async")]
+            thread_ownership_map: FxHashMap::default(),
             wrapped_failure_mt_ptr,
             #[cfg(feature = "async")]
             waker: NonNull::from(noop_waker_ref()),

src/state/raw.rs

@@ -693,6 +693,24 @@ impl RawLua {
         self.create_thread(func)
     }
 
+    /// Updates the ownership of the given implicit thread to the root user-owned thread.
+    ///
+    /// If `owner` is `None`, the thread is removed from the ownership map.
+    #[cfg(feature = "async")]
+    pub(crate) unsafe fn update_thread_ownership(&self, th: &Thread, owner: Option<*mut ffi::lua_State>) {
+        let extra = &mut *self.extra.get();
+        let th_state = th.state();
+        match owner {
+            Some(owner) => {
+                let new_owner = (extra.thread_ownership_map).get(&owner).copied().unwrap_or(owner);
+                extra.thread_ownership_map.insert(th_state, new_owner);
+            }
+            None => {
+                extra.thread_ownership_map.remove(&th_state);
+            }
+        }
+    }
+
     /// Returns the thread to the pool for later use.
     #[cfg(feature = "async")]
     pub(crate) unsafe fn recycle_thread(&self, thread: &mut Thread) {

src/thread.rs

@@ -689,6 +689,11 @@ impl<R> AsyncThread<R> {
     pub(crate) fn set_recyclable(&mut self, recyclable: bool) {
         self.recycle = recyclable;
     }
+
+    #[inline(always)]
+    pub(crate) fn thread(&self) -> &Thread {
+        &self.thread
+    }
 }
 
 #[cfg(feature = "async")]
@@ -712,6 +717,7 @@ impl<R> Drop for AsyncThread<R> {
                     if self.thread.reset_inner(status).is_ok() {
                         lua.recycle_thread(&mut self.thread);
                     }
+                    lua.update_thread_ownership(&self.thread, None);
                 }
             }
         }

tests/async.rs

@@ -7,7 +7,7 @@ use futures_util::stream::TryStreamExt;
 use tokio::sync::Mutex;
 
 use mlua::{
-    Error, Function, Lua, LuaOptions, MultiValue, ObjectLike, Result, StdLib, Table, UserData,
+    Error, Function, Lua, LuaOptions, MultiValue, ObjectLike, Result, StdLib, Table, Thread, UserData,
     UserDataMethods, UserDataRef, Value,
 };
 
@@ -718,3 +718,17 @@ fn test_async_yield_with() -> Result<()> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn test_async_current_thread() -> Result<()> {
+    let lua = Lua::new();
+
+    let get_inner_thread = lua.create_async_function(move |lua, ()| async move {
+        let f = lua.create_async_function(move |lua, ()| async move { Ok(lua.current_thread()) })?;
+        f.call_async::<Thread>(()).await
+    })?;
+    let inner_thread = get_inner_thread.call_async::<Thread>(()).await?;
+    assert_eq!(inner_thread, lua.current_thread());
+
+    Ok(())
+}