mod pool;
use std::any::Any;
use std::cell::{Cell, Ref, RefCell, RefMut};
use std::marker::PhantomData;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use rustc_hash::FxHashMap;
use crate::runtime::{SurfaceHandle, current_surface};
type TaskId = u64;
type Waker = Arc<dyn Fn() + Send + Sync>;
static TASK_WAKER: RwLock<Option<Waker>> = RwLock::new(None);
pub fn set_task_waker(wake: impl Fn() + Send + Sync + 'static) {
*TASK_WAKER.write().unwrap_or_else(|e| e.into_inner()) = Some(Arc::new(wake));
}
fn wake_loop() {
let waker = TASK_WAKER.read().unwrap_or_else(|e| e.into_inner()).clone();
if let Some(wake) = waker {
wake();
}
}
enum Message {
Item(Box<dyn Any + Send>),
End,
}
enum Callback {
Once(Box<dyn FnOnce(Box<dyn Any + Send>)>),
Stream {
item: Box<dyn FnMut(Box<dyn Any + Send>)>,
end: Box<dyn FnOnce()>,
},
}
#[derive(Default)]
struct Mailbox {
posted: Mutex<Vec<(TaskId, Message)>>,
}
impl Mailbox {
fn post(&self, id: TaskId, message: Message) {
self.posted
.lock()
.unwrap_or_else(|e| e.into_inner())
.push((id, message));
}
}
struct PendingTask {
callback: Callback,
surface: SurfaceHandle,
cancelled: Arc<AtomicBool>,
}
#[derive(Default)]
struct TaskRegistry {
next_id: TaskId,
pending: FxHashMap<TaskId, PendingTask>,
mailbox: Arc<Mailbox>,
}
struct TaskCell(Cell<*mut RefCell<TaskRegistry>>);
impl TaskCell {
fn borrow(&self) -> Ref<'_, TaskRegistry> {
unsafe { (*self.0.get()).borrow() }
}
fn borrow_mut(&self) -> RefMut<'_, TaskRegistry> {
unsafe { (*self.0.get()).borrow_mut() }
}
}
thread_local! {
static TASKS: TaskCell = TaskCell(Cell::new(
Box::into_raw(Box::new(RefCell::new(TaskRegistry::default())))
));
}
struct WorkerEnd {
mailbox: Arc<Mailbox>,
id: TaskId,
cancelled: Arc<AtomicBool>,
}
impl WorkerEnd {
fn post(&self, message: Message) {
self.mailbox.post(self.id, message);
wake_loop();
}
}
impl Drop for WorkerEnd {
fn drop(&mut self) {
self.post(Message::End);
}
}
pub struct Task {
id: TaskId,
cancelled: Arc<AtomicBool>,
_ui_thread_only: PhantomData<*const ()>,
}
impl Task {
pub fn cancel(self) {
self.cancelled.store(true, Ordering::Relaxed);
let dropped = TASKS.with(|t| t.borrow_mut().pending.remove(&self.id));
drop(dropped);
}
pub fn is_pending(&self) -> bool {
TASKS.with(|t| t.borrow().pending.contains_key(&self.id))
}
}
fn register(surface: SurfaceHandle, callback: Callback) -> (TaskId, Arc<Mailbox>, Arc<AtomicBool>) {
TASKS.with(|t| {
let mut registry = t.borrow_mut();
let id = registry.next_id;
registry.next_id += 1;
let cancelled = Arc::new(AtomicBool::new(false));
registry.pending.insert(
id,
PendingTask {
callback,
surface,
cancelled: Arc::clone(&cancelled),
},
);
(id, Arc::clone(®istry.mailbox), cancelled)
})
}
pub fn spawn_task<T, W, F>(work: W, on_done: F) -> Task
where
T: Send + 'static,
W: FnOnce() -> T + Send + 'static,
F: FnOnce(T) + 'static,
{
let callback = Callback::Once(Box::new(move |value| {
if let Ok(value) = value.downcast::<T>() {
on_done(*value);
}
}));
let (id, mailbox, cancelled) = register(current_surface(), callback);
let end = WorkerEnd {
mailbox,
id,
cancelled: Arc::clone(&cancelled),
};
pool::submit(Box::new(move || {
let value = work();
end.post(Message::Item(Box::new(value)));
}));
Task {
id,
cancelled,
_ui_thread_only: PhantomData,
}
}
pub struct Emitter<T> {
end: Arc<WorkerEnd>,
_item: PhantomData<fn(T)>,
}
impl<T: Send + 'static> Emitter<T> {
pub fn emit(&self, item: T) {
if self.is_cancelled() {
return;
}
self.end.post(Message::Item(Box::new(item)));
}
pub fn is_cancelled(&self) -> bool {
self.end.cancelled.load(Ordering::Relaxed)
}
}
impl<T> Clone for Emitter<T> {
fn clone(&self) -> Self {
Self {
end: Arc::clone(&self.end),
_item: PhantomData,
}
}
}
pub fn spawn_stream<T, W, F, E>(work: W, mut on_item: F, on_end: E) -> Task
where
T: Send + 'static,
W: FnOnce(Emitter<T>) + Send + 'static,
F: FnMut(T) + 'static,
E: FnOnce() + 'static,
{
let callback = Callback::Stream {
item: Box::new(move |value| {
if let Ok(value) = value.downcast::<T>() {
on_item(*value);
}
}),
end: Box::new(on_end),
};
let (id, mailbox, cancelled) = register(current_surface(), callback);
let emitter = Emitter::<T> {
end: Arc::new(WorkerEnd {
mailbox,
id,
cancelled: Arc::clone(&cancelled),
}),
_item: PhantomData,
};
pool::submit(Box::new(move || work(emitter)));
Task {
id,
cancelled,
_ui_thread_only: PhantomData,
}
}
pub fn drain_tasks() {
let posted = TASKS.with(|t| {
let registry = t.borrow();
let mut posted = registry
.mailbox
.posted
.lock()
.unwrap_or_else(|e| e.into_inner());
std::mem::take(&mut *posted)
});
if posted.is_empty() {
return;
}
crate::batch(|| {
for (id, message) in posted {
let Some(task) = TASKS.with(|t| t.borrow_mut().pending.remove(&id)) else {
continue; };
let value = match message {
Message::Item(value) => value,
Message::End => {
if let Callback::Stream { end, .. } = task.callback {
let _surface = task.surface.enter();
end();
}
continue;
}
};
let still_open = {
let _surface = task.surface.enter();
match task.callback {
Callback::Once(run) => {
run(value);
None
}
Callback::Stream { mut item, end } => {
item(value);
Some(Callback::Stream { item, end })
}
}
};
if let Some(callback) = still_open {
TASKS.with(|t| {
t.borrow_mut().pending.insert(
id,
PendingTask {
callback,
surface: task.surface,
cancelled: task.cancelled,
},
)
});
}
}
});
}
pub fn cancel_tasks_for(surface: SurfaceHandle) {
let cancelled: Vec<PendingTask> = TASKS.with(|t| {
let mut registry = t.borrow_mut();
let ids: Vec<TaskId> = registry
.pending
.iter()
.filter(|(_, task)| task.surface == surface)
.map(|(id, _)| *id)
.collect();
ids.iter()
.filter_map(|id| registry.pending.remove(id))
.inspect(|task| task.cancelled.store(true, Ordering::Relaxed))
.collect()
});
drop(cancelled);
}
pub fn reset_tasks() {
let flags: Vec<Arc<AtomicBool>> = TASKS.with(|t| {
t.borrow()
.pending
.values()
.map(|task| Arc::clone(&task.cancelled))
.collect()
});
for flag in flags {
flag.store(true, Ordering::Relaxed);
}
pool::shutdown_and_join();
let (pending, mailbox) = TASKS.with(|t| {
let mut registry = t.borrow_mut();
(
std::mem::take(&mut registry.pending),
std::mem::take(&mut registry.mailbox),
)
});
drop(pending);
drop(mailbox);
}
#[cfg(test)]
fn pending_task_count() -> usize {
TASKS.with(|t| t.borrow().pending.len())
}
#[cfg(test)]
mod tests {
use std::rc::Rc;
use std::sync::atomic::AtomicUsize;
use std::time::{Duration, Instant};
use super::*;
use crate::runtime::{SurfaceEnterGuard, set_current_surface, set_surface_enter_hook};
static SERIAL: Mutex<()> = Mutex::new(());
fn serial() -> std::sync::MutexGuard<'static, ()> {
SERIAL.lock().unwrap_or_else(|e| e.into_inner())
}
fn drain_until(done: impl Fn() -> bool) {
let deadline = Instant::now() + Duration::from_secs(5);
while !done() {
assert!(Instant::now() < deadline, "task never completed");
std::thread::sleep(Duration::from_millis(1));
drain_tasks();
}
}
#[test]
fn result_crosses_back_to_the_spawning_thread() {
let _serial = serial();
let got = Rc::new(std::cell::Cell::new(0i32));
let sink = Rc::clone(&got);
spawn_task(|| 6 * 7, move |v| sink.set(v));
drain_until(|| got.get() != 0);
assert_eq!(got.get(), 42);
assert_eq!(pending_task_count(), 0);
}
#[test]
fn cancel_discards_the_result() {
let _serial = serial();
let ran = Rc::new(std::cell::Cell::new(false));
let sink = Rc::clone(&ran);
let task = spawn_task(|| 1u8, move |_| sink.set(true));
task.cancel();
assert_eq!(pending_task_count(), 0);
std::thread::sleep(Duration::from_millis(20));
drain_tasks();
assert!(!ran.get());
}
#[test]
fn a_panicking_worker_abandons_its_callback() {
let _serial = serial();
let ran = Rc::new(std::cell::Cell::new(false));
let sink = Rc::clone(&ran);
spawn_task(
|| -> u8 { panic!("worker blew up") },
move |_| sink.set(true),
);
drain_until(|| pending_task_count() == 0);
assert!(!ran.get());
}
#[test]
fn completion_reenters_the_spawning_surface() {
let _serial = serial();
set_surface_enter_hook(|handle| {
let prev = set_current_surface(handle);
SurfaceEnterGuard::new(move || {
set_current_surface(prev);
})
});
let seen = Rc::new(std::cell::Cell::new(SurfaceHandle::NONE));
let sink = Rc::clone(&seen);
{
let _surface_a = SurfaceHandle(7).enter();
spawn_task(|| (), move |()| sink.set(current_surface()));
}
let _surface_b = SurfaceHandle(9).enter();
drain_until(|| !seen.get().is_none());
assert_eq!(seen.get(), SurfaceHandle(7));
}
#[test]
fn a_finishing_worker_wakes_the_loop() {
let _serial = serial();
static WAKES: AtomicUsize = AtomicUsize::new(0);
WAKES.store(0, Ordering::SeqCst);
set_task_waker(|| {
WAKES.fetch_add(1, Ordering::SeqCst);
});
spawn_task(|| (), |()| {});
drain_until(|| pending_task_count() == 0);
assert!(WAKES.load(Ordering::SeqCst) >= 1);
}
#[test]
fn a_stream_delivers_every_item_in_order_then_closes() {
let _serial = serial();
let seen: Rc<RefCell<Vec<u32>>> = Rc::new(RefCell::new(Vec::new()));
let sink = Rc::clone(&seen);
let closed = Rc::new(std::cell::Cell::new(false));
let closed_sink = Rc::clone(&closed);
spawn_stream(
|out| {
for i in 0..8u32 {
out.emit(i);
}
},
move |i| sink.borrow_mut().push(i),
move || closed_sink.set(true),
);
drain_until(|| pending_task_count() == 0);
assert_eq!(*seen.borrow(), (0..8).collect::<Vec<u32>>());
assert!(closed.get(), "on_end must run once the worker returns");
}
#[test]
fn cancelling_a_stream_stops_its_worker() {
let _serial = serial();
let seen = Rc::new(std::cell::Cell::new(0u32));
let sink = Rc::clone(&seen);
let (started_tx, started_rx) = std::sync::mpsc::channel();
static EMITTED: AtomicUsize = AtomicUsize::new(0);
EMITTED.store(0, Ordering::SeqCst);
let task = spawn_stream(
move |out| {
started_tx.send(()).ok();
while !out.is_cancelled() {
out.emit(1u32);
EMITTED.fetch_add(1, Ordering::SeqCst);
std::thread::sleep(Duration::from_millis(1));
}
},
move |i| sink.set(sink.get() + i),
|| unreachable!("on_end must not run for a cancelled stream"),
);
started_rx.recv().expect("worker never started");
task.cancel();
assert_eq!(
pending_task_count(),
0,
"cancel releases the callback at once"
);
std::thread::sleep(Duration::from_millis(60));
let settled = EMITTED.load(Ordering::SeqCst);
std::thread::sleep(Duration::from_millis(60));
assert_eq!(
settled,
EMITTED.load(Ordering::SeqCst),
"cancel must stop the stream worker, not just mute it"
);
drain_tasks();
assert_eq!(seen.get(), 0, "no item may be delivered after cancel");
}
#[test]
fn the_pool_reuses_an_idle_thread() {
let _serial = serial();
let threads: Rc<RefCell<Vec<std::thread::ThreadId>>> = Rc::new(RefCell::new(Vec::new()));
for _ in 0..2 {
let sink = Rc::clone(&threads);
spawn_task(
|| std::thread::current().id(),
move |id| sink.borrow_mut().push(id),
);
drain_until(|| pending_task_count() == 0);
}
let threads = threads.borrow();
assert_eq!(threads.len(), 2);
assert_eq!(
threads[0], threads[1],
"the second task should have reused the idle worker"
);
}
}