use std::{
cell::{Cell, RefCell, UnsafeCell},
future::Future,
pin::Pin,
ptr::NonNull,
rc::Rc,
task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};
use super::join_handle::JoinHandle;
pub const DEFAULT_READY_BUDGET: usize = 256;
struct ErasedFuture {
data: *mut (),
poll_fn: unsafe fn(*mut (), &mut Context<'_>) -> Poll<()>,
drop_fn: unsafe fn(*mut ()),
}
impl ErasedFuture {
fn new<F>(fut: F) -> Self
where
F: Future<Output = ()> + 'static,
{
let raw = Box::into_raw(Box::new(Box::pin(fut))) as *mut ();
Self {
data: raw,
poll_fn: poll_erased_future::<F>,
drop_fn: drop_erased_future::<F>,
}
}
fn poll(&mut self, cx: &mut Context<'_>) -> Poll<()> {
unsafe { (self.poll_fn)(self.data, cx) }
}
}
impl Drop for ErasedFuture {
fn drop(&mut self) {
if self.data.is_null() {
return;
}
unsafe { (self.drop_fn)(self.data) };
self.data = std::ptr::null_mut();
}
}
unsafe fn poll_erased_future<F>(raw: *mut (), cx: &mut Context<'_>) -> Poll<()>
where
F: Future<Output = ()> + 'static,
{
let fut = unsafe { &mut *(raw as *mut Pin<Box<F>>) };
fut.as_mut().poll(cx)
}
unsafe fn drop_erased_future<F>(raw: *mut ())
where
F: Future<Output = ()> + 'static,
{
drop(unsafe { Box::from_raw(raw as *mut Pin<Box<F>>) });
}
pub struct TaskQueue {
queue_head: UnsafeCell<Option<NonNull<TaskInner>>>,
queue_tail: UnsafeCell<Option<NonNull<TaskInner>>>,
}
impl Default for TaskQueue {
fn default() -> Self {
Self::new()
}
}
impl TaskQueue {
pub fn new() -> Self {
Self {
queue_head: UnsafeCell::new(None),
queue_tail: UnsafeCell::new(None),
}
}
pub fn is_empty(&self) -> bool {
unsafe { (&*self.queue_head.get()).is_none() }
}
pub fn spawn<T, F>(&self, fut: F) -> JoinHandle<T>
where
F: Future<Output = T> + 'static,
T: 'static,
{
let (jh, complete) = JoinHandle::new_local();
self.spawn_erased(ErasedFuture::new(async move {
let v = fut.await;
complete(v);
}));
jh
}
fn spawn_erased(&self, fut: ErasedFuture) {
let task = Rc::new(TaskInner {
fut: RefCell::new(Some(fut)),
exec: NonNull::from(self),
queued: Cell::new(false),
cached_waker: RefCell::new(None),
next: Cell::new(None),
});
self.enqueue_owned(task);
}
fn enqueue_owned(&self, task: Rc<TaskInner>) {
if Self::already_queued(&task) {
return;
}
self.push_queue_ptr(Self::into_ptr(task));
}
fn enqueue_ref(&self, task: &Rc<TaskInner>) {
if Self::already_queued(task) {
return;
}
self.push_queue_ptr(Self::into_ptr(task.clone()));
}
fn push_queue_ptr(&self, ptr: NonNull<TaskInner>) {
unsafe {
ptr.as_ref().next.set(None);
let tail = &mut *self.queue_tail.get();
if let Some(tail_ptr) = *tail {
tail_ptr.as_ref().next.set(Some(ptr));
} else {
*self.queue_head.get() = Some(ptr);
}
*tail = Some(ptr);
}
}
fn pop_queue_ptr(&self) -> Option<NonNull<TaskInner>> {
let head = unsafe { &mut *self.queue_head.get() };
let ptr = (*head)?;
unsafe {
let task_ref = ptr.as_ref();
*head = task_ref.next.get();
if head.is_none() {
*self.queue_tail.get() = None;
}
task_ref.next.set(None);
}
Some(ptr)
}
fn already_queued(task: &Rc<TaskInner>) -> bool {
task.queued.replace(true)
}
fn into_ptr(task: Rc<TaskInner>) -> NonNull<TaskInner> {
unsafe { NonNull::new_unchecked(Rc::into_raw(task) as *mut TaskInner) }
}
pub fn run_ready_with_budget(&self, budget: usize) -> usize {
let mut ran = 0usize;
while ran < budget {
let Some(ptr) = self.pop_queue_ptr() else {
break;
};
let task = unsafe { Rc::from_raw(ptr.as_ptr()) };
task.queued.set(false);
task.poll();
ran += 1;
}
ran
}
pub fn run_ready_all(&self) -> usize {
let mut ran = 0usize;
while let Some(ptr) = self.pop_queue_ptr() {
let task = unsafe { Rc::from_raw(ptr.as_ptr()) };
task.queued.set(false);
task.poll();
ran += 1;
}
ran
}
}
struct TaskInner {
fut: RefCell<Option<ErasedFuture>>,
exec: NonNull<TaskQueue>,
queued: Cell<bool>,
cached_waker: RefCell<Option<Waker>>,
next: Cell<Option<NonNull<Self>>>,
}
impl TaskInner {
fn poll(self: &Rc<Self>) {
let mut cached_waker = self.cached_waker.borrow_mut();
let waker = cached_waker.get_or_insert_with(|| unsafe {
Waker::from_raw(raw_waker(Rc::into_raw(self.clone()) as *const ()))
});
let mut cx = Context::from_waker(waker);
let mut borrow = self.fut.borrow_mut();
let Some(fut) = borrow.as_mut() else { return };
if fut.poll(&mut cx).is_ready() {
*borrow = None;
let _ = cached_waker.take();
}
}
fn wake(self: &Rc<Self>) {
let exec = unsafe { self.exec.as_ref() };
exec.enqueue_ref(self);
}
}
static VTABLE: RawWakerVTable =
RawWakerVTable::new(clone_waker, wake_by_val, wake_by_ref, drop_waker);
fn raw_waker(ptr: *const ()) -> RawWaker {
RawWaker::new(ptr, &VTABLE)
}
unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
let task = unsafe { Rc::<TaskInner>::from_raw(ptr as *const TaskInner) };
let cloned = task.clone();
std::mem::forget(task);
raw_waker(Rc::into_raw(cloned) as *const ())
}
unsafe fn wake_by_val(ptr: *const ()) {
let task = unsafe { Rc::<TaskInner>::from_raw(ptr as *const TaskInner) };
task.wake();
}
unsafe fn wake_by_ref(ptr: *const ()) {
let task = unsafe { Rc::<TaskInner>::from_raw(ptr as *const TaskInner) };
task.wake();
std::mem::forget(task);
}
unsafe fn drop_waker(ptr: *const ()) {
drop(unsafe { Rc::<TaskInner>::from_raw(ptr as *const TaskInner) });
}