Implement task handle awaiting

crates/lb/src/task.rs

@@ -1,5 +1,6 @@
 use crate::runtime::spawn;
 use luaffi::{cdef, metatype};
+use luajit::{LUA_MULTRET, Type};
 use std::{ffi::c_int, time::Duration};
 use tokio::{task::JoinHandle, time::sleep};
 
@@ -19,33 +20,67 @@ impl lb_tasklib {
     }
 
     pub extern "Lua" fn spawn(&self, f: function, ...) {
-        // pack the function and its arguments into a table and pass its ref to rust
+        // pack the function and its arguments into a table and pass its ref to rust.
+        //
+        // this table is used from rust-side to call the function with its args, and it's also
+        // reused to store its return values that the task handle can return when awaited. the ref
+        // is owned by the task handle and unref'ed when it's gc'ed.
         self.__spawn(__ref(__tpack(f, variadic!())))
     }
 
     extern "Lua-C" fn __spawn(&self, key: c_int) -> lb_task {
         let handle = spawn(async move |cx| {
-            // SAFETY: key is always unique, created by __ref above
+            // SAFETY: key is always unique, created by __ref above.
             let arg = unsafe { cx.new_ref_unchecked(key) };
             let mut s = cx.guard();
             s.resize(0);
-            s.push(arg);
-            let narg = s.unpack(1, 1, None) - 1; // unpack the table containing the function to call and its args
-            if let Err(err) = s.call_async(narg, 0).await {
-                drop(s);
-                cx.report_error(&err);
+            s.push(&arg);
+            let narg = s.unpack(1, 1, None) - 1; // unpack the function and its args from the table
+            debug_assert!(s.slot(2).type_of() == Type::Function);
+            match s.call_async(narg, LUA_MULTRET).await {
+                Ok(nret) => {
+                    s.pack(1, nret); // pack the return values back into the table
+                }
+                Err(err) => {
+                    drop(s);
+                    cx.report_error(&err);
+                }
             }
+            let _ = arg.into_raw(); // the original ref is owned by the task handle and unref'ed there
         });
 
-        lb_task { handle }
+        lb_task {
+            handle: Some(handle),
+            __ref: key,
+        }
     }
 }
 
 #[cdef]
 pub struct lb_task {
     #[opaque]
-    handle: JoinHandle<()>,
+    handle: Option<JoinHandle<()>>,
+    __ref: c_int,
 }
 
 #[metatype]
-impl lb_task {}
+impl lb_task {
+    pub extern "Lua" fn r#await(&self) -> many {
+        self.__await();
+        let ret = __registry[self.__ref];
+        __tunpack(ret, 1, ret.n)
+    }
+
+    async extern "Lua-C" fn __await(&mut self) {
+        if let Some(handle) = self.handle.take() {
+            handle
+                .await
+                .unwrap_or_else(|err| panic!("task handler panicked: {err}"));
+        }
+    }
+
+    #[gc]
+    extern "Lua" fn gc(&self) {
+        __unref(self.__ref);
+    }
+}

crates/luajit/src/lib.rs

@@ -9,6 +9,7 @@ use std::{
     ffi::{CStr, CString, NulError},
     fmt,
     marker::PhantomData,
+    mem::ManuallyDrop,
     ops::{Deref, DerefMut},
     os::raw::{c_char, c_int, c_void},
     pin::Pin,
@@ -35,6 +36,11 @@ pub fn url() -> &'static str {
     LUAJIT_URL.to_str().unwrap()
 }
 
+// reexport constants
+pub use luajit_sys::{
+    LUA_ENVIRONINDEX, LUA_GLOBALSINDEX, LUA_MULTRET, LUA_NOREF, LUA_REFNIL, LUA_REGISTRYINDEX,
+};
+
 /// Lua error.
 #[derive(Debug, Error)]
 #[non_exhaustive]
@@ -477,6 +483,19 @@ pub struct Ref {
     key: c_int,
 }
 
+impl Ref {
+    /// Consumes this ref and returns the original key used to create the ref.
+    ///
+    /// This key can be used to index into the registry table ([`LUA_REGISTRYINDEX`]) to retrieve
+    /// the referenced value. The key can be converted back into a ref using
+    /// [`State::new_ref_unchecked`].
+    pub fn into_raw(self) -> c_int {
+        let Self { ref mut state, key } = *ManuallyDrop::new(self);
+        unsafe { ptr::drop_in_place(state) }
+        key
+    }
+}
+
 impl Drop for Ref {
     fn drop(&mut self) {
         // SAFETY: luaL_unref is guaranteed to not fail
@@ -871,7 +890,7 @@ impl Stack {
     /// array-part, these values are **not** cleared. The number of values popped is returned, which
     /// is always equal to `n`.
     ///
-    /// This method does not invoke any metamethods.
+    /// This method does not invoke any metamethods. The table is not popped from the stack.
     ///
     /// Equivalent to `table.pack(...)`.
     ///
@@ -907,7 +926,8 @@ impl Stack {
     /// pushed at the top of the stack in ascending order. If `i > j`, then nothing is pushed.
     /// Otherwise, `j - i + 1` values are pushed, and the number of values pushed is returned.
     ///
-    /// This method does not invoke any metamethods.
+    /// This method does not invoke any metamethods. The table is not popped from the stack or
+    /// altered in any way.
     ///
     /// Equivalent to `table.unpack(list, i, j)`.
     ///