Refactor

src/lib.rs

@@ -3,9 +3,7 @@
 #![feature(context_ext)]
 
 use std::{
-    ops::DerefMut as _,
     pin::Pin,
-    ptr,
     rc::{Rc, Weak},
     task::{ContextBuilder, Poll, local_waker_fn, waker_fn},
     time::Instant,
@@ -23,10 +21,14 @@ use send_wrapper::SendWrapper;
 use slotmap::{SlotMap, new_key_type};
 use unchecked_refcell::UncheckedRefCell;
 
-use crate::op::Op;
+use crate::ring_op::RingOpHandle;
 
-pub mod op;
-mod timer;
+pub mod ops;
+mod ring_op;
+pub mod timer;
+
+#[cfg(test)]
+mod tests;
 
 struct Task {
     ready: bool,
@@ -60,8 +62,14 @@ impl Runtime {
             .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)");
+        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)])
@@ -109,6 +117,7 @@ impl Runtime {
                 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 || {
@@ -165,15 +174,15 @@ impl Runtime {
             trace!("Polling timers");
             let wait = loop {
                 match lock.timer.poll(now) {
-                    timer::Status::Empty => {
+                    timer::TimerPoll::Empty => {
                         trace!("No timers expired");
                         break None;
                     }
-                    timer::Status::Pending(deadline) => {
+                    timer::TimerPoll::Pending(deadline) => {
                         trace!("Next timer expiring in {:?}", deadline - now);
                         break Some(Timespec::from(deadline - now));
                     }
-                    timer::Status::Expired(waker) => {
+                    timer::TimerPoll::Expired(waker) => {
                         trace!("Timer expired, waking");
                         drop(lock);
                         waker.wake();
@@ -206,10 +215,11 @@ impl Runtime {
                 };
                 trace!("Processing CQE {:?}", &cqe);
                 debug_assert!(cqe.result() >= 0);
-                let user_data = ptr::with_exposed_provenance(cqe.user_data().try_into().unwrap());
-                let op: Rc<UncheckedRefCell<Op>> = unsafe { Rc::from_raw(user_data) };
                 drop(lock);
-                Op::handle(op);
+                {
+                    let handle = unsafe { RingOpHandle::from_user_data(cqe.user_data()) };
+                    handle.complete(cqe.result(), cqe.flags());
+                }
                 lock = self.0.borrow_mut();
             }
         }
@@ -220,52 +230,3 @@ impl Runtime {
         trace!("Exiting");
     }
 }
-
-#[cfg(test)]
-mod tests {
-    use std::time::Duration;
-
-    use crate::{op::no_op, timer::delay};
-
-    use super::*;
-
-    fn init() {
-        let _ = env_logger::builder().is_test(true).try_init();
-    }
-
-    #[test]
-    fn it_works() {
-        init();
-        let mut runtime = Runtime::new();
-        runtime.spawn(Box::pin(async {
-            println!("1");
-            no_op().await;
-            println!("2");
-            delay(Duration::from_secs(1)).await;
-            println!("3");
-            no_op().await;
-            println!("4");
-        }));
-        runtime.run();
-    }
-
-    #[test]
-    fn concurrency() {
-        init();
-        let mut runtime = Runtime::new();
-        for x in 0..10 {
-            runtime.spawn(Box::pin(async move {
-                println!("Starting {}", x);
-                delay(Duration::from_millis(x * 100)).await;
-                for y in 0..10 {
-                    println!("Checkin {}ms", y * 1000 + x * 100);
-                    delay(Duration::from_secs(1)).await;
-
-                    no_op().await;
-                }
-                println!("Finishing {}", x);
-            }));
-        }
-        runtime.run();
-    }
-}

src/op.rs

@@ -1,102 +0,0 @@
-use std::{
-    mem,
-    ops::DerefMut as _,
-    pin::Pin,
-    rc::Rc,
-    task::{Context, LocalWaker, Poll},
-};
-
-use io_uring::opcode::Nop;
-use log::trace;
-use unchecked_refcell::UncheckedRefCell;
-
-use crate::{RT, Runtime, RuntimeInner, op::NoOp::Cancelled};
-
-pub struct OpEntry(Rc<UncheckedRefCell<Op>>);
-
-pub enum Op {
-    NoOp(NoOp),
-}
-
-enum NoOp {
-    New,
-    Live(LocalWaker),
-    Done,
-    Cancelled,
-}
-
-pub fn no_op() -> OpEntry {
-    OpEntry(Rc::new(UncheckedRefCell::new(Op::NoOp(NoOp::New))))
-}
-
-impl Future for OpEntry {
-    type Output = ();
-
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        let mut lock = self.0.borrow_mut();
-        match &mut *lock {
-            Op::NoOp(NoOp::Cancelled) => unreachable!(),
-            Op::NoOp(NoOp::Done) => Poll::Ready(()),
-            Op::NoOp(op @ NoOp::New) => {
-                {
-                    let rt: &Runtime = cx.ext().downcast_ref().unwrap();
-                    let mut lock = rt.0.borrow_mut();
-                    let mut sq = lock.ring.submission();
-                    let entry = Nop::new().build().user_data(
-                        Rc::into_raw(self.0.clone())
-                            .expose_provenance()
-                            .try_into()
-                            .unwrap(),
-                    );
-                    trace!("Pushing SQE {:?}", &entry);
-                    unsafe { sq.push(&entry) }.expect("failed to submit no-op");
-                }
-                *op = NoOp::Live(cx.local_waker().clone());
-                Poll::Pending
-            }
-            Op::NoOp(NoOp::Live(waker)) => {
-                waker.clone_from(cx.local_waker());
-                Poll::Pending
-            }
-        }
-    }
-}
-
-impl Op {
-    pub fn handle(this: Rc<UncheckedRefCell<Self>>) {
-        // TODO don't drop for multishot accepts
-        let mut lock = this.borrow_mut();
-        match lock.deref_mut() {
-            Op::NoOp(op) => match mem::replace(op, NoOp::Done) {
-                NoOp::New | NoOp::Done => unreachable!(),
-                NoOp::Cancelled => {}
-                NoOp::Live(waker) => waker.wake(),
-            },
-        }
-    }
-}
-
-/*
-impl Drop for OpEntry {
-    fn drop(&mut self) {
-        let mut lock = self.0.borrow_mut();
-        match lock.deref_mut() {
-            Op::NoOp(op) => {
-                match mem::replace(op, NoOp::Cancelled) {
-                    NoOp::New | NoOp::Done => {}
-                    NoOp::Cancelled => unreachable!(),
-                    NoOp::Live(_) => {
-                        RT.with(|rt| {
-                            let outer = rt.borrow();
-                            let mut lock = outer.as_ref().expect("future cancelled outside runtime thread").0.borrow_mut();
-                            let mut sq = lock.ring.submission();
-                            let entry = AsyncCancel::new(self.0.as_ptr().addr().try_into().unwrap()).build();
-                            todo!()
-                        })
-                    }
-                }
-            }
-        }
-    }
-}
-*/

src/ops/mod.rs

@@ -0,0 +1 @@
+pub mod no_op;

src/ops/no_op.rs

@@ -0,0 +1,34 @@
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use io_uring::opcode::Nop;
+
+use crate::ring_op::RingOp;
+
+pub fn no_op() -> NoOp {
+    Default::default()
+}
+
+#[derive(Default)]
+pub struct NoOp {
+    ring_op: RingOp,
+}
+
+impl Future for NoOp {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self
+            .ring_op
+            .submit_and_poll(cx, |user_data| Nop::new().build().user_data(user_data))
+        {
+            Poll::Ready(complete) => {
+                debug_assert_eq!(complete.result, 0);
+                Poll::Ready(())
+            }
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}

src/ring_op.rs

@@ -0,0 +1,117 @@
+use std::{
+    mem,
+    ops::DerefMut as _,
+    ptr,
+    rc::Rc,
+    task::{Context, LocalWaker, Poll},
+};
+
+use io_uring::{opcode::AsyncCancel, squeue::Entry};
+use log::trace;
+use unchecked_refcell::UncheckedRefCell;
+
+use crate::{RT, Runtime};
+
+pub(crate) struct RingOp(RingOpHandle);
+
+pub(crate) struct RingOpHandle(Rc<UncheckedRefCell<Inner>>);
+
+enum Inner {
+    New,
+    Live(LocalWaker),
+    Complete(Complete),
+    Cancelled,
+}
+
+#[derive(Clone, Copy)]
+pub struct Complete {
+    pub result: i32,
+    pub flags: u32,
+}
+
+impl Default for RingOp {
+    fn default() -> Self {
+        Self(RingOpHandle(Rc::new(UncheckedRefCell::new(Inner::New))))
+    }
+}
+
+impl Drop for RingOp {
+    fn drop(&mut self) {
+        // TODO handle already completed before submitting cancel
+        let mut lock = self.0.0.borrow_mut();
+        match mem::replace(lock.deref_mut(), Inner::Cancelled) {
+            Inner::New | Inner::Complete(_) => {}
+            Inner::Cancelled => unreachable!(),
+            Inner::Live(_) => RT.with(|rt| {
+                let outer = rt.borrow();
+                let lock = outer
+                    .as_ref()
+                    .expect("future cancelled outside runtime thread")
+                    .0
+                    .borrow_mut();
+                let entry = AsyncCancel::new(self.0.0.as_ptr().addr().try_into().unwrap()).build();
+                trace!("Canceling sqe {:?}", &entry);
+                let mut sq = unsafe { lock.ring.submission_shared() };
+                unsafe { sq.push(&entry) }.expect("failed to submit cancel");
+                // TODO check for full
+            }),
+        }
+    }
+}
+
+impl RingOp {
+    pub fn submit_and_poll(
+        &mut self,
+        cx: &mut Context<'_>,
+        with_entry: impl FnOnce(u64) -> Entry,
+    ) -> Poll<Complete> {
+        let mut lock = self.0.0.borrow_mut();
+        match &mut *lock {
+            Inner::Cancelled => unreachable!(),
+            Inner::Complete(complete) => Poll::Ready(*complete),
+            op @ Inner::New => {
+                {
+                    let rt: &Runtime = cx.ext().downcast_ref().unwrap(); // TODO Get this way or thread local?
+                    let mut lock = rt.0.borrow_mut();
+                    let entry = with_entry(
+                        Rc::into_raw(self.0.0.clone())
+                            .expose_provenance()
+                            .try_into()
+                            .unwrap(),
+                    );
+                    trace!("Pushing SQE {:?}", &entry);
+                    let mut sq = lock.ring.submission();
+                    unsafe { sq.push(&entry) }.expect("failed to submit no-op");
+                    // TODO check for full
+                }
+                *op = Inner::Live(cx.local_waker().clone());
+                Poll::Pending
+            }
+            Inner::Live(waker) => {
+                waker.clone_from(cx.local_waker());
+                Poll::Pending
+            }
+        }
+    }
+}
+
+impl RingOpHandle {
+    pub unsafe fn from_user_data(user_data: u64) -> Self {
+        let ptr = ptr::with_exposed_provenance(user_data.try_into().unwrap());
+        Self(unsafe { Rc::from_raw(ptr) })
+    }
+
+    pub fn complete(self, result: i32, flags: u32) {
+        // TODO don't drop for multishot accepts
+        let mut lock = self.0.borrow_mut();
+
+        match mem::replace(
+            lock.deref_mut(),
+            Inner::Complete(Complete { result, flags }),
+        ) {
+            Inner::New | Inner::Complete { .. } => unreachable!(),
+            Inner::Cancelled => {}
+            Inner::Live(waker) => waker.wake(),
+        }
+    }
+}

src/tests.rs

@@ -0,0 +1,54 @@
+use std::{
+    io::{Write as _, stdout},
+    time::Duration,
+};
+
+use crate::{
+    ops::no_op::no_op,
+    timer::{delay, delay_until},
+};
+
+use super::*;
+
+fn init() {
+    let _ = env_logger::builder().is_test(true).try_init();
+}
+
+#[test]
+fn it_works() {
+    init();
+    let mut runtime = Runtime::new();
+    runtime.spawn(Box::pin(async {
+        let mut out = stdout().lock();
+        writeln!(&mut out, "1").unwrap();
+        no_op().await;
+        writeln!(&mut out, "2").unwrap();
+        delay(Duration::from_secs(1)).await;
+        writeln!(&mut out, "3").unwrap();
+        no_op().await;
+        writeln!(&mut out, "4").unwrap();
+    }));
+    runtime.run();
+}
+
+#[test]
+fn concurrency() {
+    init();
+    let mut runtime = Runtime::new();
+    let base = Instant::now();
+    for x in 0..10 {
+        runtime.spawn(Box::pin(async move {
+            println!("Starting {}", x);
+            let base = base + Duration::from_millis(x * 100);
+            delay_until(base).await;
+            for y in 0..5 {
+                println!("Checkin {}ms", y * 1000 + x * 100);
+                delay_until(base + Duration::from_secs(y + 1)).await;
+
+                no_op().await;
+            }
+            println!("Finishing {}", x);
+        }));
+    }
+    runtime.run();
+}

src/timer.rs

@@ -1,17 +1,17 @@
 use std::{
-    borrow::BorrowMut as _,
     cell::RefCell,
     cmp::Reverse,
     collections::BinaryHeap,
+    mem,
     pin::Pin,
-    rc::{Rc, Weak},
+    rc::Rc,
     task::{Context, LocalWaker, Poll},
     time::{Duration, Instant},
 };
 
 use slab::Slab;
 
-use crate::RT;
+use crate::Runtime;
 
 #[derive(PartialEq, Eq, PartialOrd, Ord)]
 struct Entry {
@@ -20,73 +20,84 @@ struct Entry {
 }
 
 enum Slot {
-    Live,
     Armed(LocalWaker),
+    Complete,
     Cancelled,
     // Potential future optimization: Delayed(Instant, T),
 }
 
 #[derive(Default)]
-pub struct TimerHeap {
+pub(crate) struct TimerHeap {
     slots: Slab<Rc<RefCell<Slot>>>,
     queue: BinaryHeap<Reverse<Entry>>,
 }
 
-pub enum Status {
+pub(crate) enum TimerPoll {
     Empty,
     Pending(Instant),
     Expired(LocalWaker),
 }
 
 impl TimerHeap {
-    pub fn insert(&mut self, deadline: Instant) -> Weak<RefCell<Slot>> {
-        let slot = Rc::new(RefCell::new(Slot::Live));
-        let weak = Rc::downgrade(&slot);
+    fn insert(&mut self, deadline: Instant, waker: LocalWaker) -> Rc<RefCell<Slot>> {
+        let slot = Rc::new(RefCell::new(Slot::Armed(waker)));
+        let handle = Rc::clone(&slot);
         let slot_id = self.slots.insert(slot);
         self.queue.push(Reverse(Entry { deadline, slot_id }));
-        weak
+        handle
     }
 
-    pub fn poll(&mut self, now: Instant) -> Status {
+    pub fn poll(&mut self, now: Instant) -> TimerPoll {
         loop {
             let slot = match self.queue.peek() {
-                None => break Status::Empty,
+                None => break TimerPoll::Empty,
                 Some(&Reverse(Entry { deadline, slot_id })) => {
                     if deadline > now {
-                        // We should check if cancelled here
-                        break Status::Pending(deadline);
+                        // We could check if cancelled here
+                        break TimerPoll::Pending(deadline);
                     }
                     self.queue.pop().unwrap();
                     self.slots.remove(slot_id)
                 }
             };
-            match Rc::into_inner(slot).unwrap().into_inner() {
-                Slot::Live => {}
-                Slot::Armed(waker) => break Status::Expired(waker),
+            let mut lock = slot.borrow_mut();
+            match mem::replace(&mut *lock, Slot::Complete) {
+                Slot::Armed(waker) => break TimerPoll::Expired(waker),
                 Slot::Cancelled => {}
+                Slot::Complete => unreachable!(),
             }
         }
     }
 }
 
-pub struct Timer {
-    slot: Weak<RefCell<Slot>>,
+pub enum Timer {
+    Pending(Instant),
+    Queued(Queued),
 }
 
+pub struct Queued(Rc<RefCell<Slot>>);
+
 impl Future for Timer {
     type Output = ();
 
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
-        match self.slot.upgrade() {
-            None => Poll::Ready(()),
-            Some(slot) => {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match &mut *self {
+            Timer::Pending(deadline) => {
+                let waker = cx.local_waker().to_owned();
+                let rt: &Runtime = cx.ext().downcast_ref().unwrap(); // TODO Get this way or thread local?
+                let mut lock = rt.0.borrow_mut();
+                *self = Timer::Queued(Queued(lock.timer.insert(*deadline, waker)));
+                Poll::Pending
+            }
+            Timer::Queued(Queued(slot)) => {
                 let mut lock = slot.as_ref().borrow_mut();
-                match &**lock.borrow_mut() {
-                    &Slot::Live | &Slot::Armed(_) => {
-                        *lock = Slot::Armed(cx.local_waker().clone());
+                match &mut *lock {
+                    Slot::Armed(old_waker) => {
+                        old_waker.clone_from(cx.local_waker());
                         Poll::Pending
                     }
-                    &Slot::Cancelled => Poll::Ready(()),
+                    Slot::Cancelled => unreachable!(),
+                    Slot::Complete => Poll::Ready(()),
                 }
             }
         }
@@ -95,20 +106,15 @@ impl Future for Timer {
 
 impl Drop for Timer {
     fn drop(&mut self) {
-        if let Some(slot) = self.slot.upgrade() {
-            *slot.as_ref().borrow_mut() = Slot::Cancelled;
+        match self {
+            Timer::Pending(_) => {}
+            Timer::Queued(Queued(slot)) => *slot.as_ref().borrow_mut() = Slot::Cancelled,
         }
     }
 }
 
 pub fn delay_until(deadline: Instant) -> Timer {
-    Timer {
-        slot: RT.with(|rt| {
-            let rt = rt.borrow();
-            let mut lock = rt.as_ref().unwrap().0.as_ref().borrow_mut();
-            lock.timer.insert(deadline)
-        }),
-    }
+    Timer::Pending(deadline)
 }
 
 pub fn delay(duration: Duration) -> Timer {