uring-async/src/lib.rs

#![feature(local_waker)]
#![feature(waker_fn)]
#![feature(context_ext)]

use std::{
    pin::Pin,
    rc::{Rc, Weak},
    task::{ContextBuilder, Poll, local_waker_fn, waker_fn},
    time::Instant,
};

use io_uring::{
    IoUring,
    opcode::{Close, Nop, Socket},
    register::Restriction,
    types::{SubmitArgs, Timespec},
};
use libc::ETIME;
use log::trace;
use send_wrapper::SendWrapper;
use slotmap::{SlotMap, new_key_type};
use unchecked_refcell::UncheckedRefCell;

use crate::ring_op::RingOpHandle;

pub mod ops;
mod ring_op;
pub mod timer;

#[cfg(test)]
mod tests;

struct Task {
    ready: bool,
    fut: Option<Pin<Box<dyn Future<Output = ()>>>>,
}

new_key_type! { struct TaskId; }

pub struct RuntimeInner {
    timer: timer::TimerHeap,
    ring: IoUring,

    tasks: SlotMap<TaskId, Task>,
    runqueue: Vec<TaskId>,
}

pub struct Runtime(Rc<UncheckedRefCell<RuntimeInner>>);
struct RuntimeWeak(Weak<UncheckedRefCell<RuntimeInner>>);

thread_local! {
    static RT: UncheckedRefCell<Option<Runtime>> = const { UncheckedRefCell::new(None) };
}

impl Runtime {
    pub fn new() -> Self {
        let ring = IoUring::builder()
            .setup_cqsize(256)
            .setup_r_disabled()
            .setup_defer_taskrun()
            .setup_single_issuer()
            .setup_no_sqarray()
            .build(128)
            .expect("failed to create ring");
        assert!(
            ring.params().is_feature_nodrop(),
            "missing IORING_FEAT_NODROP (kernel 5.19)"
        );
        assert!(
            ring.params().is_feature_ext_arg(),
            "missing IORING_FEAT_EXT_ARG (kernel 5.11)"
        );
        let submitter = ring.submitter();
        submitter
            .register_restrictions(&mut [
                Restriction::sqe_op(Nop::CODE),
                Restriction::sqe_op(Socket::CODE),
                Restriction::sqe_op(Close::CODE),
            ])
            .expect("failed to register restrictions on ring");
        submitter
            .register_enable_rings()
            .expect("failed to enable ring");
        Self(Rc::new(UncheckedRefCell::new(RuntimeInner {
            timer: Default::default(),
            ring,

            tasks: Default::default(),
            runqueue: Default::default(),
        })))
    }

    pub fn spawn(&self, fut: Pin<Box<dyn Future<Output = ()>>>) {
        let mut inner = self.0.borrow_mut();
        let task_id = inner.tasks.insert(Task {
            fut: Some(fut),
            ready: true,
        });
        inner.runqueue.push(task_id);
    }

    fn downgrade(&self) -> RuntimeWeak {
        RuntimeWeak(Rc::downgrade(&self.0))
    }

    pub fn run(&mut self) {
        RT.with(|cell| {
            let mut lock = cell.borrow_mut();
            if lock.is_some() {
                panic!("runtime already initialized on this thread");
            }
            *lock = Some(Runtime(Rc::clone(&self.0)));
        });

        let mut lock = self.0.borrow_mut();
        while !lock.tasks.is_empty() {
            // Run until no ready tasks, no timers, no future CQEs

            trace!("Polling {} tasks", lock.tasks.len());
            while let Some(task_id) = lock.runqueue.pop() {
                trace!("Polling task {:?}", task_id);
                lock.tasks[task_id].ready = false;

                // TODO better wakers
                let waker = {
                    let inner = SendWrapper::new(self.downgrade());
                    waker_fn(move || {
                        if let Some(inner) = inner.0.upgrade() {
                            trace!("Waking {:?}", task_id);
                            let mut lock = inner.borrow_mut();
                            let ready = &mut lock.tasks[task_id].ready;
                            if !*ready {
                                trace!("Enqueueing {:?}", task_id);
                                *ready = true;
                                lock.runqueue.push(task_id);
                            }
                        }
                    })
                };
                let local_waker = {
                    let inner = SendWrapper::new(self.downgrade());
                    local_waker_fn(move || {
                        if let Some(inner) = inner.0.upgrade() {
                            trace!("Local waking {:?}", task_id);
                            let mut lock = inner.borrow_mut();
                            let ready = &mut lock.tasks[task_id].ready;
                            if !*ready {
                                trace!("Local enqueueing {:?}", task_id);
                                *ready = true;
                                lock.runqueue.push(task_id);
                            }
                        }
                    })
                };

                // TODO we shouldn't need to take task out
                let mut task = lock.tasks[task_id].fut.take().unwrap();
                drop(lock);
                let mut cx = ContextBuilder::from_waker(&waker)
                    .local_waker(&local_waker)
                    .ext(self)
                    .build();
                // Unlock runtime while running user code
                let outcome = task.as_mut().poll(&mut cx);
                lock = self.0.borrow_mut();

                match outcome {
                    Poll::Ready(()) => {
                        lock.tasks.remove(task_id);
                    }
                    Poll::Pending => {
                        lock.tasks[task_id].fut = Some(task);
                    }
                }
            }

            let now = Instant::now();
            trace!("Polling timers");
            let wait = loop {
                match lock.timer.poll(now) {
                    timer::TimerPoll::Empty => {
                        trace!("No timers expired");
                        break None;
                    }
                    timer::TimerPoll::Pending(deadline) => {
                        trace!("Next timer expiring in {:?}", deadline - now);
                        break Some(Timespec::from(deadline - now));
                    }
                    timer::TimerPoll::Expired(waker) => {
                        trace!("Timer expired, waking");
                        drop(lock);
                        waker.wake();
                        lock = self.0.borrow_mut();
                    }
                }
            };

            let mut args = SubmitArgs::new();
            trace!("Entering ring");
            let want = match &wait {
                Some(wait) => {
                    args = args.timespec(wait);
                    1
                }
                None => 0,
            };
            match lock.ring.submitter().submit_with_args(want, &args) {
                Ok(_) => {}
                Err(err) if err.raw_os_error() == Some(ETIME) && wait.is_some() => {}
                e => {
                    e.expect("failed to enter ring");
                }
            }

            loop {
                let entry = lock.ring.completion().next();
                let Some(cqe) = entry else {
                    break;
                };
                trace!("Processing CQE {:?}", &cqe);
                drop(lock);
                {
                    let handle = unsafe { RingOpHandle::from_user_data(cqe.user_data()) };
                    handle.complete(cqe.result(), cqe.flags());
                }
                lock = self.0.borrow_mut();
            }
        }
        RT.with(|cell| {
            let mut lock = cell.borrow_mut();
            *lock = None;
        });
        trace!("Exiting");
    }
}