use std::cell::RefCell;
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::rc::Rc;
use std::sync::mpsc::{self, Receiver, RecvTimeoutError, Sender};
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use crate::vm::libs::background::Completion;
use crate::vm::libs::coroutine;
use crate::vm::machine::VM;
use crate::vm::value::{Value, VmTable};
enum Wake {
Elapsed(Instant),
Args(Vec<Value>),
Fail(String),
}
pub(crate) type Waker = Box<dyn Fn() + Send>;
pub(crate) struct CompletionSender {
tx: Sender<Completion>,
waker: Arc<Mutex<Option<Waker>>>,
}
impl CompletionSender {
pub(crate) fn send(self, completion: Completion) {
if self.tx.send(completion).is_ok()
&& let Some(wake) = &*self.waker.lock().unwrap_or_else(|e| e.into_inner())
{
wake();
}
}
}
pub struct Scheduler {
deferred: VecDeque<(Value, Vec<Value>)>,
sleeping: BTreeMap<(Instant, u64), (Value, Wake)>,
next_seq: u64,
ready: VecDeque<(Value, Wake)>,
waiting: HashMap<u64, Value>,
next_job: u64,
completions: Receiver<Completion>,
sender: Sender<Completion>,
waker: Arc<Mutex<Option<Waker>>>,
task_ids: HashSet<usize>,
failed: usize,
}
impl Default for Scheduler {
fn default() -> Self {
let (sender, completions) = mpsc::channel();
Scheduler {
deferred: VecDeque::new(),
sleeping: BTreeMap::new(),
next_seq: 0,
ready: VecDeque::new(),
waiting: HashMap::new(),
next_job: 0,
completions,
sender,
waker: Arc::new(Mutex::new(None)),
task_ids: HashSet::new(),
failed: 0,
}
}
}
impl Scheduler {
fn sleep_until(&mut self, at: Instant, thread: Value, wake: Wake) {
let seq = self.next_seq;
self.next_seq += 1;
self.sleeping.insert((at, seq), (thread, wake));
}
fn is_idle(&self) -> bool {
self.deferred.is_empty()
&& self.sleeping.is_empty()
&& self.ready.is_empty()
&& self.waiting.is_empty()
}
fn forget(&mut self, thread: &Rc<RefCell<VmTable>>) {
let same = |v: &Value| matches!(v, Value::Table(t) if Rc::ptr_eq(t, thread));
self.deferred.retain(|(t, _)| !same(t));
self.sleeping.retain(|_, (t, _)| !same(t));
self.ready.retain(|(t, _)| !same(t));
self.waiting.retain(|_, t| !same(t));
if let Some(id) = coroutine_id(thread) {
self.task_ids.remove(&id);
}
}
}
fn coroutine_id(thread: &Rc<RefCell<VmTable>>) -> Option<usize> {
thread.borrow().get_str("_id").as_usize()
}
pub fn create_task_lib() -> Value {
let t = Rc::new(RefCell::new(VmTable::new()));
let mut b = t.borrow_mut();
b.set_str("spawn", crate::native!("task.spawn", task_spawn));
b.set_str("defer", crate::native!("task.defer", task_defer));
b.set_str("delay", crate::native!("task.delay", task_delay));
b.set_str("wait", crate::native!("task.wait", task_wait));
b.set_str("cancel", crate::native!("task.cancel", task_cancel));
b.frozen = true;
drop(b);
Value::Table(t)
}
fn task_thread(
vm: &mut VM,
arg: Option<&Value>,
func: &str,
index: usize,
) -> Result<Value, String> {
let thread = match arg {
Some(f @ (Value::Closure(_) | Value::Native(_))) => coroutine::new_coroutine(vm, f.clone()),
Some(t @ Value::Table(_)) if t.type_name() == "thread" => t.clone(),
_ => {
return Err(format!(
"bad argument #{} to '{}' (function or thread expected)",
index, func
));
}
};
if let Value::Table(t) = &thread
&& let Some(id) = coroutine_id(t)
{
vm.tasks.task_ids.insert(id);
}
Ok(thread)
}
fn seconds_arg(args: &[Value], index: usize, func: &str) -> Result<f64, String> {
let seconds = match args.get(index) {
None | Some(Value::Nil) => 0.0,
Some(v) => crate::vm::libs::args::number(v).map_err(|_| {
format!(
"bad argument #{} to '{}' (number expected)",
index + 1,
func
)
})?,
};
if seconds.is_nan() || seconds < 0.0 {
return Err(format!(
"bad argument #{} to '{}' (seconds must be a non-negative number)",
index + 1,
func
));
}
Ok(seconds)
}
fn deadline(now: Instant, seconds: f64) -> Instant {
Duration::try_from_secs_f64(seconds)
.ok()
.and_then(|d| now.checked_add(d))
.unwrap_or_else(|| now + Duration::from_secs(60 * 60 * 24 * 365 * 100))
}
fn resume_task(vm: &mut VM, thread: Value, args: Vec<Value>) {
resume_with(vm, thread, Wake::Args(args));
}
fn resume_with(vm: &mut VM, thread: Value, wake: Wake) {
if coroutine::is_suspended(vm, &thread) {
let outcome = match wake {
Wake::Elapsed(since) => coroutine::coroutine_resume(
vm,
&[thread.clone(), Value::Float(since.elapsed().as_secs_f64())],
),
Wake::Args(args) => {
let mut resume_args = Vec::with_capacity(args.len() + 1);
resume_args.push(thread.clone());
resume_args.extend(args);
coroutine::coroutine_resume(vm, &resume_args)
}
Wake::Fail(message) => coroutine::resume_with_error(vm, &thread, message),
};
let failure = match outcome {
Ok(res) => match res.first() {
Some(Value::Bool(false)) => {
Some(res.get(1).map(|v| v.to_string()).unwrap_or_default())
}
_ => None,
},
Err(err) => Some(err),
};
if let Some(message) = failure {
eprintln!("task error: {}", message);
vm.tasks.failed += 1;
}
}
if !coroutine::is_suspended(vm, &thread)
&& let Value::Table(t) = &thread
&& let Some(id) = coroutine_id(t)
{
vm.tasks.task_ids.remove(&id);
}
}
fn task_spawn(vm: &mut VM, args: &[Value]) -> Result<Vec<Value>, String> {
let thread = task_thread(vm, args.first(), "task.spawn", 1)?;
resume_task(vm, thread.clone(), args[1..].to_vec());
Ok(vec![thread])
}
fn task_defer(vm: &mut VM, args: &[Value]) -> Result<Vec<Value>, String> {
let thread = task_thread(vm, args.first(), "task.defer", 1)?;
vm.tasks
.deferred
.push_back((thread.clone(), args[1..].to_vec()));
Ok(vec![thread])
}
fn task_delay(vm: &mut VM, args: &[Value]) -> Result<Vec<Value>, String> {
let seconds = seconds_arg(args, 0, "task.delay")?;
let thread = task_thread(vm, args.get(1), "task.delay", 2)?;
let rest = args.get(2..).unwrap_or_default().to_vec();
let at = deadline(Instant::now(), seconds);
vm.tasks.sleep_until(at, thread.clone(), Wake::Args(rest));
Ok(vec![thread])
}
fn task_wait(vm: &mut VM, args: &[Value]) -> Result<Vec<Value>, String> {
let seconds = seconds_arg(args, 0, "task.wait")?;
let start = Instant::now();
let at = deadline(start, seconds);
if let Some(handle) = coroutine::yieldable_handle(vm) {
vm.tasks
.sleep_until(at, Value::Table(handle), Wake::Elapsed(start));
return coroutine::coroutine_yield(vm, &[]);
}
run_until(vm, Some(at));
Ok(vec![Value::Float(start.elapsed().as_secs_f64())])
}
fn task_cancel(vm: &mut VM, args: &[Value]) -> Result<Vec<Value>, String> {
let Some(Value::Table(thread)) = args.first().filter(|v| v.type_name() == "thread") else {
return Err("bad argument #1 to 'task.cancel' (thread expected)".to_string());
};
vm.tasks.forget(thread);
match coroutine::coroutine_close(vm, &args[..1])?.as_slice() {
[Value::Bool(false), message, ..] => Err(message.to_string()),
_ => Ok(vec![]),
}
}
pub(crate) fn set_waker(vm: &mut VM, waker: Waker) {
*vm.tasks.waker.lock().unwrap_or_else(|e| e.into_inner()) = Some(waker);
}
pub(crate) fn has_waker(vm: &VM) -> bool {
vm.tasks
.waker
.lock()
.unwrap_or_else(|e| e.into_inner())
.is_some()
}
pub(crate) fn current_task(vm: &VM) -> Option<Value> {
let id = vm.current_co?;
if !vm.tasks.task_ids.contains(&id) {
return None;
}
coroutine::yieldable_handle(vm).map(Value::Table)
}
pub(crate) fn expect_completion(vm: &mut VM, thread: Value) -> (u64, CompletionSender) {
let id = vm.tasks.next_job;
vm.tasks.next_job += 1;
vm.tasks.waiting.insert(id, thread);
let sender = CompletionSender {
tx: vm.tasks.sender.clone(),
waker: vm.tasks.waker.clone(),
};
(id, sender)
}
fn file_completion(vm: &mut VM, completion: Completion) {
let Some(thread) = vm.tasks.waiting.remove(&completion.id) else {
return;
};
let wake = match (completion.finish)(vm) {
Ok(values) => Wake::Args(values),
Err(err) => Wake::Fail(err),
};
vm.tasks.ready.push_back((thread, wake));
}
fn drain_completions(vm: &mut VM) {
while let Ok(completion) = vm.tasks.completions.try_recv() {
file_completion(vm, completion);
}
}
fn run_pass(vm: &mut VM, budget: Option<Duration>) -> bool {
let start = Instant::now();
drain_completions(vm);
let deferred = std::mem::take(&mut vm.tasks.deferred);
vm.tasks.ready.extend(
deferred
.into_iter()
.map(|(thread, args)| (thread, Wake::Args(args))),
);
while let Some(entry) = vm.tasks.sleeping.first_entry() {
if entry.key().0 > start {
break;
}
let due = entry.remove();
vm.tasks.ready.push_back(due);
}
let due = vm.tasks.ready.len();
let mut ran = false;
for _ in 0..due {
if ran && budget.is_some_and(|b| start.elapsed() >= b) {
break;
}
let Some((thread, wake)) = vm.tasks.ready.pop_front() else {
break;
};
resume_with(vm, thread, wake);
ran = true;
}
ran
}
fn run_until(vm: &mut VM, until: Option<Instant>) {
loop {
let ran = run_pass(vm, None);
if until.is_some_and(|u| u <= Instant::now()) {
return;
}
if ran {
continue;
}
let next = vm.tasks.sleeping.first_key_value().map(|(k, _)| k.0);
let wake = match (next, until) {
(Some(n), Some(u)) => Some(n.min(u)),
(Some(n), None) => Some(n),
(None, Some(u)) => Some(u),
(None, None) if vm.tasks.waiting.is_empty() => return,
(None, None) => None,
};
if vm.tasks.waiting.is_empty() {
let now = Instant::now();
if let Some(wake) = wake
&& wake > now
{
std::thread::sleep(wake - now);
}
continue;
}
let received = match wake {
Some(wake) => vm
.tasks
.completions
.recv_timeout(wake.saturating_duration_since(Instant::now())),
None => vm
.tasks
.completions
.recv()
.map_err(|_| RecvTimeoutError::Disconnected),
};
if let Ok(completion) = received {
file_completion(vm, completion);
}
}
}
pub fn run_until_idle(vm: &mut VM) -> Result<(), String> {
run_until(vm, None);
debug_assert!(vm.tasks.is_idle());
match std::mem::take(&mut vm.tasks.failed) {
0 => Ok(()),
1 => Err("a task raised an error".to_string()),
n => Err(format!("{} tasks raised an error", n)),
}
}
pub fn run_due(vm: &mut VM) {
run_pass(vm, None);
vm.tasks.failed = 0;
}
pub(crate) fn run_for(vm: &mut VM, budget: Duration) {
run_pass(vm, Some(budget));
}
const COMPLETION_POLL: Duration = Duration::from_millis(20);
pub fn next_due(vm: &mut VM) -> Option<Instant> {
drain_completions(vm);
let now = Instant::now();
if !vm.tasks.deferred.is_empty() || !vm.tasks.ready.is_empty() {
return Some(now);
}
let sleeper = vm.tasks.sleeping.first_key_value().map(|(k, _)| k.0);
if vm.tasks.waiting.is_empty() || has_waker(vm) {
return sleeper;
}
let poll = now + COMPLETION_POLL;
Some(sleeper.map_or(poll, |s| s.min(poll)))
}
#[cfg(test)]
mod tests {
use crate::vm::machine::VM;
use crate::vm::value::Value;
fn run_program(vm: &mut VM, src: &str) -> Result<Value, String> {
let proto = crate::compiler::compile_program(src, "test", crate::compiler::Options::RUN)?;
vm.execute_program(proto)
}
fn run_native_program(vm: &mut VM, src: &str) -> Result<Value, String> {
let proto =
crate::compiler::compile_program(src, "test", crate::compiler::Options::UNCHECKED)?;
vm.execute_program(proto)
}
fn global(vm: &VM, name: &str) -> Value {
vm.globals
.get(&crate::vm::value::StrRef::from(name))
.cloned()
.unwrap_or(Value::Nil)
}
#[test]
fn test_tasks_run_after_the_main_chunk() {
let mut vm = VM::new();
run_program(
&mut vm,
"_G.ticks = 0\n\
task.spawn(function()\n\
\u{20} while task.wait(0.001) do\n\
\u{20} _G.ticks += 1\n\
\u{20} if _G.ticks == 3 then break end\n\
\u{20} end\n\
end)\n\
task.delay(0.002, function(v: string) _G.delayed = v end, \"done\")\n\
_G.main_saw = _G.ticks",
)
.expect("program failed");
assert_eq!(global(&vm, "main_saw").to_string(), "0");
assert_eq!(global(&vm, "ticks").to_string(), "3");
assert_eq!(global(&vm, "delayed").to_string(), "done");
assert!(vm.tasks.is_idle());
}
#[test]
fn test_task_errors_fail_the_program_after_the_rest_ran() {
let mut vm = VM::new();
let err = run_program(
&mut vm,
"task.spawn(function() error(\"boom\") end)\n\
task.delay(0.001, function() error(\"again\") end)\n\
task.delay(0.002, function() _G.ran = true end)\n\
_G.main_done = true",
)
.expect_err("a failing task must fail the program");
assert_eq!(err, "2 tasks raised an error");
assert!(matches!(global(&vm, "main_done"), Value::Bool(true)));
assert!(matches!(global(&vm, "ran"), Value::Bool(true)));
}
#[test]
fn test_budgeted_passes_leave_the_rest_for_later() {
let mut vm = VM::new();
vm.execute(
crate::compiler::compile_program(
"_G.count = 0\n\
for i = 1, 5 do\n\
\u{20} task.defer(function()\n\
\u{20} local t = os.clock()\n\
\u{20} while os.clock() - t < 0.01 do end\n\
\u{20} _G.count += 1\n\
\u{20} end)\n\
end",
"test",
crate::compiler::Options::RUN,
)
.expect("compiles"),
)
.expect("runs");
super::run_for(&mut vm, std::time::Duration::from_millis(1));
assert_eq!(global(&vm, "count").to_string(), "1");
super::run_until_idle(&mut vm).expect("the rest run");
assert_eq!(global(&vm, "count").to_string(), "5");
}
#[test]
fn test_background_work_resumes_tasks() {
fn slow_double(vm: &mut VM, args: &[Value]) -> Result<Vec<Value>, String> {
let n = crate::vm::libs::args::number_arg(args, 0, "n")?;
crate::vm::libs::background::run(
vm,
move || {
std::thread::sleep(std::time::Duration::from_millis(100));
if n < 0.0 {
Err("negative".to_string())
} else {
Ok(n * 2.0)
}
},
|_vm, doubled| Ok(vec![Value::Float(doubled)]),
)
}
let mut vm = VM::new();
vm.register_native(crate::native_def!("slow_double", slow_double));
let start = std::time::Instant::now();
let err = run_native_program(
&mut vm,
"_G.sum = 0\n\
for i = 1, 5 do\n\
\u{20} task.spawn(function() local d = slow_double(i) _G.sum += d end)\n\
end\n\
task.spawn(function() slow_double(-1) _G.after = true end)\n\
_G.direct = slow_double(21)",
)
.expect_err("the failing task fails the program");
assert_eq!(err, "a task raised an error");
assert_eq!(global(&vm, "direct").to_string(), "42");
assert_eq!(global(&vm, "sum").to_string(), "30");
assert!(matches!(global(&vm, "after"), Value::Nil));
assert!(start.elapsed() < std::time::Duration::from_millis(500));
assert!(vm.tasks.is_idle());
}
#[test]
fn test_cancel_drops_background_results() {
fn slow(vm: &mut VM, _args: &[Value]) -> Result<Vec<Value>, String> {
crate::vm::libs::background::run(
vm,
|| {
std::thread::sleep(std::time::Duration::from_millis(50));
Ok(())
},
|_vm, ()| Ok(vec![Value::Bool(true)]),
)
}
let mut vm = VM::new();
vm.register_native(crate::native_def!("slow", slow));
run_native_program(
&mut vm,
"local t = task.spawn(function() _G.got = slow() end)\n\
task.cancel(t)\n\
task.wait(0.15)",
)
.expect("program runs");
assert!(matches!(global(&vm, "got"), Value::Nil));
assert!(vm.tasks.is_idle());
}
#[test]
fn test_plain_coroutines_do_not_yield_for_background_work() {
fn slow(vm: &mut VM, _args: &[Value]) -> Result<Vec<Value>, String> {
crate::vm::libs::background::run(vm, || Ok(7), |_vm, n| Ok(vec![Value::Int(n)]))
}
let mut vm = VM::new();
vm.register_native(crate::native_def!("slow", slow));
run_native_program(
&mut vm,
"local gen = coroutine.wrap(function() coroutine.yield(slow()) end)\n\
_G.got = gen()",
)
.expect("program runs");
assert_eq!(global(&vm, "got").to_string(), "7");
}
}