use crate::async_engine::AsyncEngine;
use crate::io_uring::buffer_ring::BufferRingError;
use crate::io_uring::debounce::Debouncer;
use crate::io_uring::ops::accept::AcceptTask;
use crate::io_uring::ops::async_cancel::AsyncCancelTask;
use crate::io_uring::ops::connect::ConnectTask;
use crate::io_uring::ops::futex::FutexWaitTask;
use crate::io_uring::ops::futex::FutexWakeTask;
use crate::io_uring::ops::poll::PollTask;
use crate::io_uring::ops::poll_external::PollExternalTask;
use crate::io_uring::ops::read_multishot::ReadMultishotTask;
use crate::io_uring::ops::read_write::ReadWriteTask;
use crate::io_uring::ops::read_write_no_cancel::ReadWriteNoCancelTask;
pub use crate::io_uring::ops::read_write_vec::WriteVecCache;
use crate::io_uring::ops::read_write_vec::WriteVecTask;
use crate::io_uring::ops::recvmsg::RecvmsgTask;
use crate::io_uring::ops::sendmsg::SendmsgTask;
use crate::io_uring::ops::timeout::TimeoutTask;
use crate::io_uring::ops::timeout_external::TimeoutExternalTask;
use crate::io_uring::ops::timeout_link::TimeoutLinkTask;
pub use crate::io_uring::ops::writev_external::WritevData;
use crate::io_uring::ops::writev_external::WritevExternalTask;
use crate::io_uring::pending_result::PendingResults;
use crate::io_uring::sys::IORING_CQE_F_MORE;
use crate::io_uring::sys::IORING_ENTER_GETEVENTS;
use crate::io_uring::sys::IORING_FEAT_NODROP;
use crate::io_uring::sys::IORING_OFF_CQ_RING;
use crate::io_uring::sys::IORING_OFF_SQ_RING;
use crate::io_uring::sys::IORING_OFF_SQES;
use crate::io_uring::sys::IORING_SETUP_COOP_TASKRUN;
use crate::io_uring::sys::IORING_SETUP_DEFER_TASKRUN;
use crate::io_uring::sys::IORING_SETUP_SINGLE_ISSUER;
use crate::io_uring::sys::IORING_SETUP_SUBMIT_ALL;
use crate::io_uring::sys::IOSQE_IO_LINK;
use crate::io_uring::sys::io_uring_cqe;
use crate::io_uring::sys::io_uring_enter;
use crate::io_uring::sys::io_uring_params;
use crate::io_uring::sys::io_uring_setup;
use crate::io_uring::sys::io_uring_sqe;
use crate::utils::asyncevent::AsyncEvent;
use crate::utils::bitfield::Bitfield;
use crate::utils::bitflags::BitflagsExt;
use crate::utils::buf::Buf;
use crate::utils::copyhashmap::FCopyHashMap;
use crate::utils::errorfmt::ErrorFmt;
use crate::utils::numcell::NumCell;
use crate::utils::ptr_ext::MutPtrExt;
use crate::utils::ptr_ext::PtrExt;
use crate::utils::stack::Stack;
use crate::utils::syncqueue::SyncQueue;
use jay_algorithms::mmap::Mmapped;
use jay_algorithms::mmap::mmap;
use jay_algorithms::oserror::OsError;
pub use ops::TaskResultExt;
pub use ops::futex::FutexObj;
pub use ops::poll_external::PendingPoll;
pub use ops::poll_external::PollCallback;
pub use ops::read_multishot::PendingReadMultishot;
pub use ops::read_multishot::ReadMultishotCallback;
pub use ops::timeout_external::PendingTimeout;
pub use ops::timeout_external::TimeoutCallback;
use std::cell::Cell;
use std::cell::RefCell;
use std::cell::UnsafeCell;
use std::rc::Rc;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering::Acquire;
use std::sync::atomic::Ordering::Relaxed;
use std::sync::atomic::Ordering::Release;
use std::task::Waker;
use thiserror::Error;
use uapi::OwnedFd;
use uapi::c;
macro_rules! map_err {
($n:expr) => {{
let n = $n;
if n < 0 {
Err(jay_algorithms::oserror::OsError::from(-n as uapi::c::c_int))
} else {
Ok(n)
}
}};
}
pub mod buffer_ring;
mod debounce;
mod ops;
mod pending_result;
mod sys;
#[derive(Debug, Error)]
pub enum IoUringError {
#[error(transparent)]
OsError(#[from] OsError),
#[error("Could not create an io-uring")]
CreateUring(#[source] OsError),
#[error("The kernel does not support the IORING_FEAT_NODROP feature")]
NoDrop,
#[error("Could not map the submission queue ring")]
MapSqRing(#[source] OsError),
#[error("Could not map the submission queue entries")]
MapSqEntries(#[source] OsError),
#[error("Could not map the completion queue ring")]
MapCqRing(#[source] OsError),
#[error("The io-uring has already been destroyed")]
Destroyed,
#[error("io_uring_enter failed")]
Enter(#[source] OsError),
#[error("Kernel sent invalid cmsg data")]
InvalidCmsgData,
#[error(transparent)]
BufferRing(#[from] BufferRingError),
}
pub struct IoUring {
ring: Rc<IoUringData>,
}
impl Drop for IoUring {
fn drop(&mut self) {
self.ring.kill();
}
}
impl IoUring {
pub fn new(eng: &Rc<AsyncEngine>, entries: u32) -> Result<Rc<Self>, IoUringError> {
let feature_levels = [
IORING_SETUP_SUBMIT_ALL, IORING_SETUP_COOP_TASKRUN, IORING_SETUP_SINGLE_ISSUER, IORING_SETUP_DEFER_TASKRUN, ];
let mut feature_levels = &feature_levels[..];
let mut params;
let fd = loop {
params = io_uring_params::default();
for &flags in feature_levels {
params.flags |= flags;
}
match io_uring_setup(entries, &mut params) {
Ok(f) => break f,
Err(e) => {
if let Some((_, levels)) = feature_levels.split_last() {
feature_levels = levels;
} else {
return Err(IoUringError::CreateUring(e));
}
}
}
};
if !params.features.contains(IORING_FEAT_NODROP) {
return Err(IoUringError::NoDrop);
}
let sqmap_map = mmap(
(params.sq_off.array + params.sq_entries * 4) as _,
c::PROT_READ | c::PROT_WRITE,
c::MAP_SHARED | c::MAP_POPULATE,
fd.raw(),
IORING_OFF_SQ_RING as _,
);
let sqmap_map = match sqmap_map {
Ok(map) => map,
Err(e) => return Err(IoUringError::MapSqRing(e)),
};
let sqesmap_map = mmap(
params.sq_entries as usize * size_of::<io_uring_sqe>(),
c::PROT_READ | c::PROT_WRITE,
c::MAP_SHARED | c::MAP_POPULATE,
fd.raw(),
IORING_OFF_SQES as _,
);
let sqesmap_map = match sqesmap_map {
Ok(map) => map,
Err(e) => return Err(IoUringError::MapSqEntries(e)),
};
let cqmap_map = mmap(
params.cq_off.cqes as usize + params.cq_entries as usize * size_of::<io_uring_cqe>(),
c::PROT_READ | c::PROT_WRITE,
c::MAP_SHARED | c::MAP_POPULATE,
fd.raw(),
IORING_OFF_CQ_RING as _,
);
let cqmap_map = match cqmap_map {
Ok(map) => map,
Err(e) => return Err(IoUringError::MapCqRing(e)),
};
let sqmask = unsafe {
(sqmap_map.ptr as *const u8)
.add(params.sq_off.ring_mask as _)
.cast::<u32>()
.read()
};
let sqhead = unsafe {
(sqmap_map.ptr as *const u8)
.add(params.sq_off.head as _)
.cast()
};
let sqtail = unsafe {
(sqmap_map.ptr as *const u8)
.add(params.sq_off.tail as _)
.cast()
};
let sqmap = unsafe {
let base = (sqmap_map.ptr as *const u8)
.add(params.sq_off.array as _)
.cast();
std::slice::from_raw_parts(base, params.sq_entries as _)
};
let sqesmap = unsafe {
let base = (sqesmap_map.ptr as *const u8).cast();
std::slice::from_raw_parts(base, params.sq_entries as _)
};
let cqmask = unsafe {
(cqmap_map.ptr as *const u8)
.add(params.cq_off.ring_mask as _)
.cast::<u32>()
.read()
};
let cqhead = unsafe {
(cqmap_map.ptr as *const u8)
.add(params.cq_off.head as _)
.cast()
};
let cqtail = unsafe {
(cqmap_map.ptr as *const u8)
.add(params.cq_off.tail as _)
.cast()
};
let cqmap = unsafe {
let base = (cqmap_map.ptr as *const u8)
.add(params.cq_off.cqes as _)
.cast();
std::slice::from_raw_parts(base, params.cq_entries as _)
};
let data = Rc::new(IoUringData {
destroyed: Cell::new(false),
fd,
eng: eng.clone(),
_sqesmap_map: sqesmap_map,
_sqmap_map: sqmap_map,
sqmask,
sqlen: params.sq_entries,
sqhead,
sqtail,
sqmap,
sqesmap,
_cqmap_map: cqmap_map,
cqmask,
cqhead,
cqtail,
cqmap,
cqes_consumed: Default::default(),
next: Default::default(),
to_encode: Default::default(),
pending_in_kernel: Default::default(),
tasks: Default::default(),
pending_results: Default::default(),
cached_read_writes: Default::default(),
cached_read_writes_no_cancel: Default::default(),
cached_cancels: Default::default(),
cached_polls: Default::default(),
cached_polls_external: Default::default(),
cached_sendmsg: Default::default(),
cached_recvmsg: Default::default(),
cached_timeouts: Default::default(),
cached_timeouts_external: Default::default(),
cached_timeout_links: Default::default(),
cached_cmsg_bufs: Default::default(),
cached_connects: Default::default(),
cached_accepts: Default::default(),
cached_futex_wakes: Default::default(),
cached_futex_waits: Default::default(),
cached_write_vecs: Default::default(),
cached_read_multishot: Default::default(),
cached_writev_external: Default::default(),
fd_ids_scratch: Default::default(),
iteration: Default::default(),
yields: Default::default(),
multishot_tasks: Default::default(),
buffer_ring_ids: Default::default(),
});
Ok(Rc::new(Self { ring: data }))
}
pub fn stop(&self) {
self.ring.kill();
}
pub fn run(&self) -> Result<(), IoUringError> {
let res = self.ring.run();
self.ring.kill();
res
}
pub fn cancel(&self, id: IoUringTaskId) {
self.ring.cancel_task(id);
}
pub fn debouncer(&self, max: u64) -> Debouncer {
Debouncer {
cur: Default::default(),
max,
iteration: Cell::new(self.ring.iteration.get()),
ring: self.ring.clone(),
}
}
}
struct IoUringData {
destroyed: Cell<bool>,
fd: OwnedFd,
eng: Rc<AsyncEngine>,
_sqesmap_map: Mmapped,
_sqmap_map: Mmapped,
sqmask: u32,
sqlen: u32,
sqhead: *const AtomicU32,
sqtail: *const AtomicU32,
sqmap: *const [Cell<c::c_uint>],
sqesmap: *const [UnsafeCell<io_uring_sqe>],
_cqmap_map: Mmapped,
cqmask: u32,
cqhead: *const AtomicU32,
cqtail: *const AtomicU32,
cqmap: *const [Cell<io_uring_cqe>],
cqes_consumed: AsyncEvent,
next: IoUringTaskIds,
to_encode: SyncQueue<IoUringTaskId>,
pending_in_kernel: FCopyHashMap<IoUringTaskId, ()>,
tasks: FCopyHashMap<IoUringTaskId, TaskPlus>,
multishot_tasks: FCopyHashMap<IoUringTaskId, Rc<dyn MultishotTask>>,
pending_results: PendingResults,
cached_read_writes: Stack<Box<ReadWriteTask>>,
cached_read_writes_no_cancel: Stack<Box<ReadWriteNoCancelTask>>,
cached_cancels: Stack<Box<AsyncCancelTask>>,
cached_polls: Stack<Box<PollTask>>,
cached_polls_external: Stack<Box<PollExternalTask>>,
cached_sendmsg: Stack<Box<SendmsgTask>>,
cached_recvmsg: Stack<Box<RecvmsgTask>>,
cached_timeouts: Stack<Box<TimeoutTask>>,
cached_timeouts_external: Stack<Box<TimeoutExternalTask>>,
cached_timeout_links: Stack<Box<TimeoutLinkTask>>,
cached_cmsg_bufs: Stack<Buf>,
cached_connects: Stack<Box<ConnectTask>>,
cached_accepts: Stack<Box<AcceptTask>>,
cached_futex_wakes: Stack<Box<FutexWakeTask>>,
cached_futex_waits: Stack<Box<FutexWaitTask>>,
cached_write_vecs: Stack<Box<WriteVecTask>>,
cached_read_multishot: Stack<Box<ReadMultishotTask>>,
cached_writev_external: Stack<Box<WritevExternalTask>>,
fd_ids_scratch: RefCell<Vec<c::c_int>>,
iteration: NumCell<u64>,
yields: SyncQueue<Waker>,
buffer_ring_ids: RefCell<Bitfield>,
}
struct TaskPlus {
task: Box<dyn Task>,
has_timeout: bool,
multishot: bool,
}
unsafe trait Task: 'static {
fn id(&self) -> IoUringTaskId;
fn complete(self: Box<Self>, ring: &IoUringData, cqe: &io_uring_cqe);
fn encode(&self, sqe: &mut io_uring_sqe);
fn is_cancel(&self) -> bool {
false
}
fn has_timeout(&self) -> bool {
false
}
fn multishot(&self) -> Option<Rc<dyn MultishotTask>> {
None
}
}
unsafe trait MultishotTask: 'static {
fn complete_partial(self: Rc<Self>, cqe: &io_uring_cqe);
}
impl IoUringData {
fn run(&self) -> Result<(), IoUringError> {
let mut to_submit = 0;
loop {
self.iteration.fetch_add(1);
while let Some(ev) = self.yields.pop() {
ev.wake();
}
loop {
self.eng.dispatch();
if self.destroyed.get() {
return Ok(());
}
if !self.dispatch_completions() {
break;
}
}
to_submit += self.encode();
let res = {
let (to_submit, mut min_complete, flags) = if to_submit == 0 {
(0, 1, IORING_ENTER_GETEVENTS)
} else if self.to_encode.is_empty() {
(to_submit as _, 1, IORING_ENTER_GETEVENTS)
} else {
(!0, 0, 0)
};
if self.yields.is_not_empty() {
min_complete = 0;
}
io_uring_enter(self.fd.raw(), to_submit, min_complete, flags)
};
let mut submitted_any = false;
match res {
Ok(n) => {
if n > 0 {
submitted_any = true;
}
to_submit -= n;
}
Err(e) => {
if not_matches!(e.0, c::EAGAIN | c::EBUSY | c::EINTR) {
return Err(IoUringError::Enter(e));
}
}
}
if to_submit > 0 && !submitted_any {
let res = io_uring_enter(self.fd.raw(), 0, 1, IORING_ENTER_GETEVENTS);
if let Err(e) = res {
if e.0 != c::EINTR {
return Err(IoUringError::Enter(e));
}
}
}
}
}
fn dispatch_completions(&self) -> bool {
unsafe {
let mut head = self.cqhead.deref().load(Relaxed);
let tail = self.cqtail.deref().load(Acquire);
if head == tail {
return false;
}
while head != tail {
let idx = (head & self.cqmask) as usize;
let entry = self.cqmap.deref()[idx].get();
head = head.wrapping_add(1);
self.cqhead.deref().store(head, Release);
let id = IoUringTaskId(entry.user_data);
if entry.flags.contains(IORING_CQE_F_MORE) {
if let Some(pending) = self.multishot_tasks.get(&id) {
pending.complete_partial(&entry);
}
} else {
if let Some(pending) = self.tasks.remove(&id) {
if pending.multishot {
self.multishot_tasks.remove(&id);
}
self.pending_in_kernel.remove(&id);
pending.task.complete(self, &entry);
}
}
}
self.cqhead.deref().store(head, Release);
self.cqes_consumed.trigger();
true
}
}
fn encode(&self) -> usize {
let tasks = self.tasks.lock();
let mut encoded = 0;
unsafe {
let mut tail = self.sqtail.deref().load(Relaxed);
let head = self.sqhead.deref().load(Acquire);
let available = self.sqlen - tail.wrapping_sub(head);
while encoded < available {
let Some(id) = self.to_encode.pop() else {
break;
};
let Some(task) = tasks.get(&id) else { continue };
let has_timeout = task.has_timeout;
if has_timeout && (available - encoded) < 2 {
self.to_encode.push_front(id);
break;
}
self.pending_in_kernel.set(id, ());
let idx = (tail & self.sqmask) as usize;
let sqe = self.sqesmap.deref()[idx].get().deref_mut();
self.sqmap.deref()[idx].set(idx as _);
*sqe = Default::default();
sqe.user_data = id.raw();
task.task.encode(sqe);
if has_timeout {
sqe.flags |= IOSQE_IO_LINK;
}
tail = tail.wrapping_add(1);
encoded += 1;
}
self.sqtail.deref().store(tail, Release);
}
encoded as usize
}
fn id(&self) -> Cancellable<'_> {
Cancellable {
id: self.id_raw(),
data: self,
}
}
fn id_raw(&self) -> IoUringTaskId {
self.next.next()
}
fn cancel_task(&self, id: IoUringTaskId) {
if !self.tasks.contains(&id) {
return;
}
if !self.pending_in_kernel.contains(&id) {
let cqe = io_uring_cqe {
user_data: id.raw(),
res: -c::ECANCELED,
flags: 0,
};
let task = self.tasks.remove(&id).unwrap();
if task.multishot {
self.multishot_tasks.remove(&id);
}
task.task.complete(self, &cqe);
return;
}
self.cancel_task_in_kernel(id);
}
fn schedule(&self, t: Box<impl Task>) {
self.schedule_(t.id(), t.has_timeout(), t.multishot(), t);
}
fn schedule_(
&self,
id: IoUringTaskId,
has_timeout: bool,
m: Option<Rc<dyn MultishotTask>>,
t: Box<dyn Task>,
) {
assert!(!self.destroyed.get());
let mut multishot = false;
if let Some(m) = m {
self.multishot_tasks.set(id, m);
multishot = true;
}
self.to_encode.push(id);
self.tasks.set(
id,
TaskPlus {
task: t,
has_timeout,
multishot,
},
);
}
fn check_destroyed(&self) -> Result<(), IoUringError> {
if self.destroyed.get() {
Err(IoUringError::Destroyed)
} else {
Ok(())
}
}
fn kill(&self) {
self.eng.stop();
let mut to_cancel = vec![];
for task in self.tasks.lock().values() {
if !task.task.is_cancel() {
to_cancel.push(task.task.id());
}
}
for task in to_cancel {
self.cancel_task(task);
}
self.destroyed.set(true);
while !self.tasks.is_empty() {
self.encode();
let _ = io_uring_enter(self.fd.raw(), u32::MAX, 0, 0);
let res = io_uring_enter(self.fd.raw(), 0, 1, IORING_ENTER_GETEVENTS);
if let Err(e) = res {
panic!("Could not wait for io_uring to drain: {}", ErrorFmt(e));
}
while self.dispatch_completions() {
}
}
}
fn cmsg_buf(&self) -> Buf {
self.cached_cmsg_bufs
.pop()
.unwrap_or_else(|| Buf::new(1024))
}
}
linear_ids!(IoUringTaskIds, IoUringTaskId, u64);
#[allow(clippy::derivable_impls)]
impl Default for IoUringTaskId {
fn default() -> Self {
Self(0)
}
}
struct Cancellable<'a> {
id: IoUringTaskId,
data: &'a IoUringData,
}
impl Drop for Cancellable<'_> {
fn drop(&mut self) {
self.data.cancel_task(self.id);
}
}