uring-async/src/ring_op.rs

use std::{
    io, mem,
    ops::DerefMut as _,
    pin::Pin,
    ptr,
    rc::Rc,
    task::{Context, LocalWaker, Poll},
};

use io_uring::{opcode::AsyncCancel, squeue::Entry};
use log::trace;
use pin_project::pin_project;
use unchecked_refcell::UncheckedRefCell;

use crate::{RT, Runtime};

#[derive(Default)]
pub(crate) struct RingOp(RingOpHandle);

#[derive(Default, Clone)]
pub(crate) struct RingOpHandle(Rc<UncheckedRefCell<Inner>>);

#[derive(Default)]
enum Inner {
    #[default]
    New,
    Live(LocalWaker),
    Complete(Complete),
    Cancelled,
}

#[derive(Clone, Copy)]
pub(crate) struct Complete {
    pub result: i32,
    pub _flags: u32,
}

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 poll_submit(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        with_entry: impl FnOnce() -> 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().user_data(self.0.clone().into_user_data());
                    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 fn into_user_data(self) -> u64 {
        Rc::into_raw(self.0).expose_provenance().try_into().unwrap()
    }

    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 self for multishot accepts
        let mut lock = self.0.borrow_mut();

        match mem::replace(
            lock.deref_mut(),
            Inner::Complete(Complete {
                result,
                _flags: flags,
            }),
        ) {
            Inner::New | Inner::Complete { .. } => unreachable!(),
            Inner::Cancelled => {}
            Inner::Live(waker) => waker.wake(),
        }
    }
}

pub fn submit_and_forget(entry: Entry) {
    RT.with(|rt| {
        let outer = rt.borrow();
        let lock = outer
            .as_ref()
            .expect("socket closed outside runtime thread")
            .0
            .borrow_mut();

        // Start cancelled to ignore result
        let handle = RingOpHandle(Rc::new(UncheckedRefCell::new(Inner::Cancelled)));
        let entry = entry.user_data(handle.into_user_data());
        trace!("Submitting and forgetting {:?}", &entry);
        let mut sq = unsafe { lock.ring.submission_shared() };
        unsafe { sq.push(&entry) }.expect("failed to submit cancel");
        // TODO check for full
    });
}

#[pin_project]
struct RingOpFut<E: FnMut() -> Entry> {
    #[pin]
    ring_op: RingOp,
    with_entry: E,
}

impl<E: FnMut() -> Entry> Future for RingOpFut<E> {
    type Output = Complete;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.project();
        this.ring_op.poll_submit(cx, this.with_entry)
    }
}

pub async fn ring_op(with_entry: impl FnMut() -> Entry) -> Complete {
    RingOpFut {
        ring_op: Default::default(),
        with_entry,
    }
    .await
}

pub async fn ring_op_result(with_entry: impl FnMut() -> Entry) -> io::Result<i32> {
    let Complete { result, .. } = ring_op(with_entry).await;
    if result < 0 {
        Err(io::Error::from_raw_os_error(-result))
    } else {
        Ok(result)
    }
}