diff --git a/src/lib.rs b/src/lib.rs index 314dc66..7748e7a 100644 --- a/src/lib.rs +++ b/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> = 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(); - } -} diff --git a/src/op.rs b/src/op.rs deleted file mode 100644 index 2e0c5e6..0000000 --- a/src/op.rs +++ /dev/null @@ -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>); - -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 { - 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>) { - // 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!() - }) - } - } - } - } - } -} -*/ diff --git a/src/ops/mod.rs b/src/ops/mod.rs new file mode 100644 index 0000000..7fe2eaf --- /dev/null +++ b/src/ops/mod.rs @@ -0,0 +1 @@ +pub mod no_op; diff --git a/src/ops/no_op.rs b/src/ops/no_op.rs new file mode 100644 index 0000000..e47c2e2 --- /dev/null +++ b/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 { + 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, + } + } +} diff --git a/src/ring_op.rs b/src/ring_op.rs new file mode 100644 index 0000000..6c224ac --- /dev/null +++ b/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>); + +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 { + 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(), + } + } +} diff --git a/src/tests.rs b/src/tests.rs new file mode 100644 index 0000000..4b4c6ce --- /dev/null +++ b/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(); +} diff --git a/src/timer.rs b/src/timer.rs index 93e2e06..dc303c3 100644 --- a/src/timer.rs +++ b/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>>, queue: BinaryHeap>, } -pub enum Status { +pub(crate) enum TimerPoll { Empty, Pending(Instant), Expired(LocalWaker), } impl TimerHeap { - pub fn insert(&mut self, deadline: Instant) -> Weak> { - let slot = Rc::new(RefCell::new(Slot::Live)); - let weak = Rc::downgrade(&slot); + fn insert(&mut self, deadline: Instant, waker: LocalWaker) -> Rc> { + 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>, +pub enum Timer { + Pending(Instant), + Queued(Queued), } +pub struct Queued(Rc>); + impl Future for Timer { type Output = (); - fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - match self.slot.upgrade() { - None => Poll::Ready(()), - Some(slot) => { + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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 {