dope-kqueue 0.2.1

Thin io_uring adaptor with "Manifolds"
Documentation
use std::cell::RefCell;
use std::io;
use std::os::fd::RawFd;
use std::pin::Pin;
use std::ptr::NonNull;
use std::task::{Context, Poll, Waker};

use crate::Errno;

use super::Driver;
use super::udata::Udata;

pub(crate) const TAG_REQUEST: u8 = 0xFF;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct RequestId {
    pub(crate) index: u16,
    pub(crate) epoch: u32,
}

impl RequestId {
    pub(crate) fn new(index: u16, epoch: u32) -> Self {
        Self { index, epoch }
    }

    pub(crate) fn into_udata(self) -> Udata {
        Udata::pack(TAG_REQUEST, self.index, self.epoch)
    }
}

pub(crate) struct RequestSlot {
    pub(crate) used: bool,
    pub(crate) epoch: u32,
    pub(crate) fd: RawFd,
    pub(crate) filter: i16,
    pub(crate) op: OneShotOp,
    pub(crate) outcome: Option<io::Result<u32>>,
    pub(crate) waker: Option<Waker>,
    pub(crate) armed: bool,
}

#[derive(Clone, Copy, Debug)]
pub(crate) enum OneShotOp {
    Connect,
    Send { ptr: *const u8, len: u32 },
    SendMsg { msg: *const libc::msghdr },
    Wake,
}

impl OneShotOp {
    pub(crate) fn run(self, fd: RawFd) -> io::Result<u32> {
        match self {
            Self::Connect => crate::driver::Sock::check_so_error(fd).map(|_| 0u32),
            Self::Send { ptr, len } => {
                let n = unsafe { libc::send(fd, ptr.cast(), len as usize, 0) };
                Errno::syscall_result(n)
            }
            Self::SendMsg { msg } => {
                let n = unsafe { libc::sendmsg(fd, msg, 0) };
                Errno::syscall_result(n)
            }
            Self::Wake => Ok(0),
        }
    }
}

#[derive(Default)]
pub(crate) struct Requests {
    slots: RefCell<Vec<RequestSlot>>,
    free: RefCell<Vec<u16>>,
}

impl Requests {
    pub(crate) fn alloc(
        &self,
        fd: RawFd,
        filter: i16,
        op: OneShotOp,
        outcome: Option<io::Result<u32>>,
    ) -> RequestId {
        let mut slots = self.slots.borrow_mut();
        let mut free = self.free.borrow_mut();
        let (index, epoch) = if let Some(idx) = free.pop() {
            let slot = &mut slots[idx as usize];
            slot.used = true;
            slot.fd = fd;
            slot.filter = filter;
            slot.op = op;
            slot.outcome = outcome;
            slot.waker = None;
            slot.armed = false;
            (idx, slot.epoch)
        } else {
            let idx = u16::try_from(slots.len())
                .expect("dope-kqueue invariant violated: request slot index exceeds u16");
            slots.push(RequestSlot {
                used: true,
                epoch: 1,
                fd,
                filter,
                op,
                outcome,
                waker: None,
                armed: false,
            });
            (idx, 1)
        };
        RequestId::new(index, epoch)
    }

    pub(crate) fn with_mut<R>(
        &self,
        id: RequestId,
        f: impl FnOnce(&mut RequestSlot) -> R,
    ) -> Option<R> {
        let mut slots = self.slots.borrow_mut();
        let slot = slots.get_mut(id.index as usize)?;
        if !slot.used || slot.epoch != id.epoch {
            return None;
        }
        Some(f(slot))
    }

    pub(crate) fn release(&self, id: RequestId) {
        let mut slots = self.slots.borrow_mut();
        let mut free = self.free.borrow_mut();
        let Some(slot) = slots.get_mut(id.index as usize) else {
            return;
        };
        if !slot.used || slot.epoch != id.epoch {
            return;
        }
        slot.used = false;
        slot.epoch = slot.epoch.wrapping_add(1).max(1);
        slot.outcome = None;
        slot.waker = None;
        slot.armed = false;
        free.push(id.index);
    }

    pub(crate) fn fire(&self, id: RequestId, result: io::Result<u32>) -> Option<Waker> {
        self.with_mut(id, |slot| {
            slot.outcome = Some(result);
            slot.armed = false;
            slot.waker.take()
        })?
    }
}

enum RequestState<T> {
    Pending { data: Box<T>, id: RequestId },
    Done,
}

pub struct Request<T: 'static> {
    state: RequestState<T>,
    driver: NonNull<Driver>,
}

impl<T: 'static> Request<T> {
    pub(crate) fn new(driver: &mut Driver, data: Box<T>, id: RequestId) -> Self {
        Self {
            state: RequestState::Pending { data, id },
            driver: NonNull::from(driver),
        }
    }
}

impl<T: Unpin + 'static> Request<T> {
    pub(crate) async fn await_or_err(submit: Result<Self, (io::Error, T)>) -> (io::Result<u32>, T) {
        match submit {
            Ok(req) => {
                let (res, _, op) = req.await;
                (res, op)
            }
            Err((e, op)) => (Err(e), op),
        }
    }
}

impl<T: Unpin + 'static> std::future::Future for Request<T> {
    type Output = (io::Result<u32>, u32, T);

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.get_mut();
        let id = match &me.state {
            RequestState::Pending { id, .. } => *id,
            RequestState::Done => panic!("dope-kqueue: polled completed Request"),
        };

        let driver = unsafe { me.driver.as_mut() };
        let outcome = driver
            .requests
            .with_mut(id, |slot| {
                if slot.outcome.is_some() {
                    slot.outcome.take()
                } else {
                    slot.waker = Some(cx.waker().clone());
                    None
                }
            })
            .flatten();
        match outcome {
            Some(result) => {
                let RequestState::Pending { data, id } =
                    std::mem::replace(&mut me.state, RequestState::Done)
                else {
                    unreachable!()
                };
                driver.requests.release(id);
                Poll::Ready((result, 0, *data))
            }
            None => Poll::Pending,
        }
    }
}

impl<T: 'static> Drop for Request<T> {
    fn drop(&mut self) {
        let RequestState::Pending { id, .. } =
            std::mem::replace(&mut self.state, RequestState::Done)
        else {
            return;
        };

        let driver = unsafe { self.driver.as_mut() };
        let cleanup = driver
            .requests
            .with_mut(id, |slot| (slot.fd, slot.filter, slot.armed));
        if let Some((fd, filter, true)) = cleanup {
            driver.disarm_filter(fd, filter);
        }
        driver.requests.release(id);
    }
}