dope_runtime/executor/

task_queue.rs

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) });
}