use std::collections::VecDeque;
use std::sync::{Condvar, Mutex, OnceLock};
use std::time::Duration;
use crate::vm::machine::VM;
use crate::vm::value::Value;
pub(crate) const MAX_WORKERS: usize = 16;
const IDLE_EXIT: Duration = Duration::from_secs(30);
pub(crate) type Finish = Box<dyn FnOnce(&mut VM) -> Result<Vec<Value>, String> + Send>;
pub(crate) struct Completion {
pub id: u64,
pub finish: Finish,
}
type Job = Box<dyn FnOnce() + Send>;
#[derive(Default)]
struct Queue {
jobs: VecDeque<Job>,
workers: usize,
idle: usize,
}
struct Pool {
queue: Mutex<Queue>,
wake: Condvar,
}
fn pool() -> &'static Pool {
static POOL: OnceLock<Pool> = OnceLock::new();
POOL.get_or_init(|| Pool {
queue: Mutex::new(Queue::default()),
wake: Condvar::new(),
})
}
fn submit(job: Job) {
let pool = pool();
let mut queue = pool.queue.lock().unwrap_or_else(|e| e.into_inner());
queue.jobs.push_back(job);
if queue.jobs.len() > queue.idle && queue.workers < MAX_WORKERS {
let spawned = std::thread::Builder::new()
.name("neyuki-worker".to_string())
.spawn(worker);
match spawned {
Ok(_) => queue.workers += 1,
Err(_) if queue.workers == 0 => {
let job = queue.jobs.pop_back().expect("the job was just queued");
drop(queue);
job();
return;
}
Err(_) => {}
}
}
pool.wake.notify_one();
}
fn worker() {
let pool = pool();
loop {
let job = {
let mut queue = pool.queue.lock().unwrap_or_else(|e| e.into_inner());
loop {
if let Some(job) = queue.jobs.pop_front() {
break job;
}
queue.idle += 1;
let (guard, timeout) = pool
.wake
.wait_timeout(queue, IDLE_EXIT)
.unwrap_or_else(|e| e.into_inner());
queue = guard;
queue.idle -= 1;
if timeout.timed_out() && queue.jobs.is_empty() {
queue.workers -= 1;
return;
}
}
};
let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(job));
}
}
pub(crate) fn run<T, W, F>(vm: &mut VM, work: W, finish: F) -> Result<Vec<Value>, String>
where
T: Send + 'static,
W: FnOnce() -> Result<T, String> + Send + 'static,
F: FnOnce(&mut VM, T) -> Result<Vec<Value>, String> + Send + 'static,
{
let Some(thread) = crate::vm::libs::task::current_task(vm) else {
let value = work()?;
return finish(vm, value);
};
let (id, sender) = crate::vm::libs::task::expect_completion(vm, thread);
submit(Box::new(move || {
let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(work))
.unwrap_or_else(|_| Err("background work panicked".to_string()));
let finish: Finish = Box::new(move |vm| finish(vm, outcome?));
sender.send(Completion { id, finish });
}));
crate::vm::libs::coroutine::coroutine_yield(vm, &[])
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::mpsc;
#[test]
fn test_pool_runs_jobs_in_parallel() {
let (tx, rx) = mpsc::channel();
let start = std::time::Instant::now();
for i in 0..8 {
let tx = tx.clone();
submit(Box::new(move || {
std::thread::sleep(Duration::from_millis(200));
let _ = tx.send(i);
}));
}
let mut seen: Vec<i32> = (0..8).map(|_| rx.recv().expect("job ran")).collect();
seen.sort();
assert_eq!(seen, (0..8).collect::<Vec<_>>());
assert!(start.elapsed() < Duration::from_millis(1200));
}
#[test]
fn test_pool_survives_a_panicking_job() {
submit(Box::new(|| panic!("job failed")));
let (tx, rx) = mpsc::channel();
submit(Box::new(move || {
let _ = tx.send(1);
}));
assert_eq!(rx.recv_timeout(Duration::from_secs(5)), Ok(1));
}
}