luby/crates/lb/src/task.rs

104 lines
3.1 KiB
Rust

//! Task library
//!
//! The `lb:task` library primitives for asynchronous communication between tasks via message
//! passing channels.
//!
//! ## Exports
//!
//! See [`lb_tasklib`] for items exported by this library.
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};
/// Items exported by the `lb:task` library.
///
/// This library can be acquired by calling
/// [`require("lb:task")`](https://www.lua.org/manual/5.1/manual.html#pdf-require) in Lua.
///
/// ```lua
/// local task = require("lb:task");
/// ```
#[cdef(module = "lb:task")]
pub struct lb_tasklib;
#[metatype]
#[include("task.lua")]
impl lb_tasklib {
#[new]
extern "Lua-C" fn new() -> Self {
Self
}
pub async extern "Lua-C" fn sleep(&self, ms: f64) {
sleep(Duration::from_secs_f64(ms / 1000.)).await;
}
pub extern "Lua" fn spawn(&self, 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.
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.
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 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: Some(handle),
__ref: key,
}
}
}
/// Handle for an asynchronous task created by [`spawn`](lb_tasklib::spawn).
#[cdef]
pub struct lb_task {
#[opaque]
handle: Option<JoinHandle<()>>,
__ref: c_int,
}
#[metatype]
impl lb_task {
pub async 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);
}
}