139 lines
3.5 KiB
Rust
139 lines
3.5 KiB
Rust
#![allow(unused)]
|
|
|
|
use {
|
|
alloc::{boxed::Box, collections::BTreeMap, sync::Arc, task::Wake},
|
|
core::{
|
|
future::Future,
|
|
pin::Pin,
|
|
task::{Context, Poll, Waker},
|
|
},
|
|
crossbeam_queue::SegQueue,
|
|
kiam::when,
|
|
slab::Slab,
|
|
spin::RwLock,
|
|
};
|
|
|
|
static SPAWN_QUEUE: RwLock<Option<SpawnQueue>> = RwLock::new(None);
|
|
pub fn spawn(future: impl Future<Output = ()> + Send + 'static) {
|
|
match &*SPAWN_QUEUE.read() {
|
|
Some(s) => s.push(Task::new(future)),
|
|
None => panic!("no task executor is running"),
|
|
}
|
|
}
|
|
|
|
pub fn yield_now() -> impl Future<Output = ()> {
|
|
struct YieldNow(bool);
|
|
impl Future for YieldNow {
|
|
type Output = ();
|
|
|
|
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
|
|
if self.0 {
|
|
Poll::Ready(())
|
|
} else {
|
|
self.0 = true;
|
|
cx.waker().wake_by_ref();
|
|
Poll::Pending
|
|
}
|
|
}
|
|
}
|
|
|
|
YieldNow(false)
|
|
}
|
|
|
|
#[derive(Default)]
|
|
pub struct Executor {
|
|
tasks: Slab<Task>,
|
|
queue: TaskQueue,
|
|
to_spawn: SpawnQueue,
|
|
wakers: BTreeMap<TaskId, Waker>,
|
|
}
|
|
|
|
impl Executor {
|
|
pub fn spawn(&mut self, future: impl Future<Output = ()> + Send + 'static) {
|
|
self.queue
|
|
.push(TaskId(self.tasks.insert(Task::new(future))));
|
|
}
|
|
|
|
pub fn run(&mut self) {
|
|
{
|
|
let mut global_spawner = SPAWN_QUEUE.write();
|
|
if global_spawner.is_some() {
|
|
panic!("Task executor is already running");
|
|
}
|
|
|
|
*global_spawner = Some(Arc::clone(&self.to_spawn));
|
|
}
|
|
|
|
loop {
|
|
when! {
|
|
let Some(id) = self
|
|
.to_spawn
|
|
.pop()
|
|
.map(|t| TaskId(self.tasks.insert(t)))
|
|
.or_else(|| self.queue.pop())
|
|
=> {
|
|
let Some(task) = self.tasks.get_mut(id.0) else {
|
|
panic!("Attempted to get task from empty slot: {}", id.0);
|
|
};
|
|
|
|
let mut cx = Context::from_waker(self.wakers.entry(id).or_insert_with(|| {
|
|
Waker::from(Arc::new(TaskWaker {
|
|
id,
|
|
queue: Arc::clone(&self.queue),
|
|
}))
|
|
}));
|
|
|
|
match task.poll(&mut cx) {
|
|
Poll::Ready(()) => {
|
|
self.tasks.remove(id.0);
|
|
self.wakers.remove(&id);
|
|
}
|
|
Poll::Pending => (),
|
|
}
|
|
},
|
|
self.tasks.is_empty() => break,
|
|
_ => (),
|
|
}
|
|
}
|
|
|
|
*SPAWN_QUEUE.write() = None;
|
|
}
|
|
}
|
|
|
|
struct Task {
|
|
future: Pin<Box<dyn Future<Output = ()> + Send>>,
|
|
}
|
|
|
|
impl Task {
|
|
pub fn new(future: impl Future<Output = ()> + Send + 'static) -> Self {
|
|
Self {
|
|
future: Box::pin(future),
|
|
}
|
|
}
|
|
|
|
fn poll(&mut self, cx: &mut Context) -> Poll<()> {
|
|
self.future.as_mut().poll(cx)
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
|
struct TaskId(usize);
|
|
|
|
type TaskQueue = Arc<SegQueue<TaskId>>;
|
|
type SpawnQueue = Arc<SegQueue<Task>>;
|
|
|
|
struct TaskWaker {
|
|
id: TaskId,
|
|
queue: TaskQueue,
|
|
}
|
|
|
|
impl Wake for TaskWaker {
|
|
fn wake(self: Arc<Self>) {
|
|
self.wake_by_ref();
|
|
}
|
|
|
|
fn wake_by_ref(self: &Arc<Self>) {
|
|
self.queue.push(self.id);
|
|
}
|
|
}
|