Fix race condition in task state table unref

crates/lb/src/task.rs

@@ -12,7 +12,7 @@ use luaffi::{
     marker::{function, many},
     metatype,
 };
-use luajit::{LUA_MULTRET, Type};
+use luajit::LUA_MULTRET;
 use std::{cell::RefCell, ffi::c_int, time::Duration};
 use tokio::{task::JoinHandle, time::sleep};
 
@@ -42,38 +42,41 @@ impl lb_tasklib {
     pub extern "Lua" fn spawn(f: function, ...) -> lb_task {
         // 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.
+        // this "state" 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.
         assert(
             r#type(f) == "function",
             concat!("function expected in argument 'f', got ", r#type(f)),
         );
-        Self::__spawn(__ref(__tpack(f, variadic!())))
+        // we need two refs: one for the spawn call, and the other for the task handle. this is to
+        // ensure the task handle isn't gc'ed and the state table unref'ed before the spawn callback
+        // runs and puts the state table on the stack.
+        let state = __tpack(f, variadic!());
+        Self::__spawn(__ref(state), __ref(state))
     }
 
-    extern "Lua-C" fn __spawn(key: c_int) -> lb_task {
+    extern "Lua-C" fn __spawn(spawn_ref: c_int, handle_ref: c_int) -> lb_task {
         let handle = spawn(async move |cx| {
-            // SAFETY: key is always unique, created by __ref above.
-            let arg = unsafe { cx.new_ref_unchecked(key) };
+            // SAFETY: handle_ref is always unique, created in Self::spawn above.
+            let state = unsafe { cx.new_ref_unchecked(spawn_ref) };
             let mut s = cx.guard();
             s.resize(0);
-            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);
+            s.push(state); // this drops the state table ref, but the table is still on the stack
+            let narg = s.unpack(1, 1, None) - 1; // unpack the function and its args from the state table
             match s.call_async(narg, LUA_MULTRET).await {
                 Ok(nret) => {
-                    s.pack(1, nret); // pack the return values back into the table
+                    s.pack(1, nret); // pack the return values back into the state 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::new(handle, key)
+        // spawn_ref is owned by the task handle and unref'ed there when the handle gets gc'ed
+        lb_task::new(handle, handle_ref)
     }
 }
 
@@ -103,8 +106,8 @@ impl lb_task {
     async extern "Lua-C" fn __await(&self) {
         if let Some(handle) = self.handle.borrow_mut().take() {
             handle
-                .await
-                .unwrap_or_else(|err| panic!("task handler panicked: {err}"));
+                .await // task handler should never panic
+                .unwrap_or_else(|err| std::panic::resume_unwind(err.into_panic()));
         }
     }
 

crates/luaffi/src/lib.lua

@@ -8,16 +8,15 @@ local function __ref(value)
   if ref ~= nil and ref ~= 0 then
     __registry[FREELIST_REF] = __registry[ref]
   else
-    ref = #__registry + 1
+    ref = rawlen(__registry) + 1
   end
   __registry[ref] = value
   return ref
 end
 
 local function __unref(ref)
-  if ref < 0 then return nil end
-  local value = __registry[ref]
-  __registry[ref] = __registry[FREELIST_REF]
-  __registry[FREELIST_REF] = ref
-  return value
+  if ref > 0 then
+    __registry[ref] = __registry[FREELIST_REF]
+    __registry[FREELIST_REF] = ref
+  end
 end