use std::collections::VecDeque;
use std::future::Future;
use std::pin::Pin;
use std::sync::{Arc, Mutex, MutexGuard};
use std::task::{Context, Poll};
use async_task::Runnable;
const DRAIN_BUDGET: usize = 100_000;
type ReadyQueue = Arc<Mutex<VecDeque<Runnable>>>;
pub(crate) struct ForegroundExecutor {
ready: ReadyQueue,
}
impl ForegroundExecutor {
pub(crate) fn new() -> ForegroundExecutor {
ForegroundExecutor {
ready: Arc::default(),
}
}
pub(crate) fn spawn<R: 'static>(&self, future: impl Future<Output = R> + 'static) -> Task<R> {
spawn_on(&self.ready, future)
}
pub(crate) fn handle(&self) -> ExecutorHandle {
ExecutorHandle {
ready: Arc::clone(&self.ready),
}
}
}
impl Drop for ForegroundExecutor {
fn drop(&mut self) {
loop {
let queued: Vec<_> = lock(&self.ready).drain(..).collect();
if queued.is_empty() {
break;
}
drop(queued);
}
}
}
#[derive(Clone)]
pub(crate) struct ExecutorHandle {
ready: ReadyQueue,
}
impl ExecutorHandle {
pub(crate) fn spawn<R: 'static>(&self, future: impl Future<Output = R> + 'static) -> Task<R> {
spawn_on(&self.ready, future)
}
pub(crate) fn drain(&self) -> usize {
let mut polled = 0usize;
loop {
let next = lock(&self.ready).pop_front();
let Some(runnable) = next else {
break;
};
polled += 1;
assert!(
polled <= DRAIN_BUDGET,
"tasks did not converge after {DRAIN_BUDGET} polls; are two tasks waking each \
other in a cycle?"
);
runnable.run();
}
polled
}
}
fn spawn_on<R: 'static>(ready: &ReadyQueue, future: impl Future<Output = R> + 'static) -> Task<R> {
let ready = Arc::clone(ready);
let schedule = move |runnable| lock(&ready).push_back(runnable);
let (runnable, task) = async_task::spawn_local(future, schedule);
runnable.schedule();
Task::spawned(task)
}
fn lock(ready: &ReadyQueue) -> MutexGuard<'_, VecDeque<Runnable>> {
ready
.lock()
.expect("the ready queue is never poisoned: nothing runs under its lock")
}
pub struct Task<T> {
inner: TaskInner<T>,
}
enum TaskInner<T> {
Ready(Option<T>),
Spawned(Option<async_task::Task<T>>),
}
impl<T> Task<T> {
pub fn ready(value: T) -> Task<T> {
Task {
inner: TaskInner::Ready(Some(value)),
}
}
pub(crate) fn spawned(task: async_task::Task<T>) -> Task<T> {
Task {
inner: TaskInner::Spawned(Some(task)),
}
}
}
impl<T> Unpin for Task<T> {}
impl<T> Future for Task<T> {
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
match &mut self.get_mut().inner {
TaskInner::Ready(value) => Poll::Ready(
value
.take()
.expect("a Task is polled until it completes, then never again"),
),
TaskInner::Spawned(task) => {
let task = task
.as_mut()
.expect("a Task is polled until it completes, then never again");
Pin::new(task).poll(cx)
}
}
}
}
impl<T> Drop for Task<T> {
fn drop(&mut self) {
if let TaskInner::Spawned(task) = &mut self.inner
&& let Some(task) = task.take()
{
task.detach();
}
}
}