kossLAN/wry
1
0
Fork 0
forked from wry/wry
Code Pull requests Activity

io_uring: move runtime into workspace crate

Browse source
This commit is contained in:
kossLAN 2026-05-29 09:22:59 -04:00
parent 03d3876888
commit c3b17db151
No known key found for this signature in database
22 changed files with 662 additions and 617 deletions
Download patch file Download diff file Expand all files Collapse all files

33
io-uring/src/debounce.rs Normal file
View file

@ -0,0 +1,33 @@
use {
crate::IoUringData,
jay_utils::numcell::NumCell,
std::{cell::Cell, future::poll_fn, rc::Rc, task::Poll},
};
pub struct Debouncer {
pub(super) cur: NumCell<u64>,
pub(super) max: u64,
pub(super) iteration: Cell<u64>,
pub(super) ring: Rc<IoUringData>,
}
impl Debouncer {
pub async fn debounce(&self) {
let iteration = self.ring.iteration.get();
if self.iteration.replace(iteration) != iteration {
self.cur.set(0);
}
if self.cur.fetch_add(1) > self.max {
poll_fn(|ctx| {
if self.ring.iteration.get() > iteration {
Poll::Ready(())
} else {
self.ring.yields.push(ctx.waker().clone());
Poll::Pending
}
})
.await;
self.cur.set(0);
}
}
}

587
io-uring/src/lib.rs Normal file
View file

@ -0,0 +1,587 @@
pub use ops::{
TaskResultExt,
poll_external::{PendingPoll, PollCallback},
timeout_external::{PendingTimeout, TimeoutCallback},
};
use {
crate::{
debounce::Debouncer,
ops::{
accept::AcceptTask, async_cancel::AsyncCancelTask, connect::ConnectTask,
poll::PollTask, poll_external::PollExternalTask, read_write::ReadWriteTask,
read_write_no_cancel::ReadWriteNoCancelTask, recvmsg::RecvmsgTask,
sendmsg::SendmsgTask, timeout::TimeoutTask, timeout_external::TimeoutExternalTask,
timeout_link::TimeoutLinkTask,
},
pending_result::PendingResults,
sys::{
IORING_ENTER_GETEVENTS, IORING_FEAT_NODROP, IORING_OFF_CQ_RING, IORING_OFF_SQ_RING,
IORING_OFF_SQES, IORING_SETUP_COOP_TASKRUN, IORING_SETUP_DEFER_TASKRUN,
IORING_SETUP_SINGLE_ISSUER, IORING_SETUP_SUBMIT_ALL, IOSQE_IO_LINK, io_uring_cqe,
io_uring_enter, io_uring_params, io_uring_setup, io_uring_sqe,
},
},
jay_async_engine::AsyncEngine,
jay_utils::{
asyncevent::AsyncEvent,
bitflags::BitflagsExt,
buf::Buf,
copyhashmap::CopyHashMap,
errorfmt::ErrorFmt,
mmap::{Mmapped, mmap},
numcell::NumCell,
oserror::OsError,
ptr_ext::{MutPtrExt, PtrExt},
stack::Stack,
syncqueue::SyncQueue,
},
std::{
cell::{Cell, RefCell, UnsafeCell},
rc::Rc,
sync::atomic::{
AtomicU32,
Ordering::{Acquire, Relaxed, Release},
},
task::Waker,
},
thiserror::Error,
uapi::{
OwnedFd,
c::{self},
},
};
macro_rules! map_err {
($n:expr) => {{
let n = $n;
if n < 0 {
Err(jay_utils::oserror::OsError::from(-n as uapi::c::c_int))
} else {
Ok(n)
}
}};
}
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,
}
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, // 5.18
IORING_SETUP_COOP_TASKRUN, // 5.19
IORING_SETUP_SINGLE_ISSUER, // 6.0
IORING_SETUP_DEFER_TASKRUN, // 6.1
];
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()
};
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()
};
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(),
fd_ids_scratch: Default::default(),
iteration: Default::default(),
yields: 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: CopyHashMap<IoUringTaskId, ()>,
tasks: CopyHashMap<IoUringTaskId, Box<dyn Task>>,
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>>,
fd_ids_scratch: RefCell<Vec<c::c_int>>,
iteration: NumCell<u64>,
yields: SyncQueue<Waker>,
}
unsafe trait Task {
fn id(&self) -> IoUringTaskId;
fn complete(self: Box<Self>, ring: &IoUringData, res: i32);
fn encode(&self, sqe: &mut io_uring_sqe);
fn is_cancel(&self) -> bool {
false
}
fn has_timeout(&self) -> bool {
false
}
}
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 !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 let Some(pending) = self.tasks.remove(&id) {
self.pending_in_kernel.remove(&id);
pending.complete(self, entry.res);
}
}
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 id = match self.to_encode.pop() {
Some(t) => t,
_ => break,
};
let task = match tasks.get(&id) {
Some(t) => t,
_ => 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.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) {
self.tasks
.remove(&id)
.unwrap()
.complete(self, -c::ECANCELED);
return;
}
self.cancel_task_in_kernel(id);
}
fn schedule(&self, t: Box<dyn Task>) {
assert!(!self.destroyed.get());
self.to_encode.push(t.id());
self.tasks.set(t.id(), t);
}
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.is_cancel() {
to_cancel.push(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() {
// nothing
}
}
}
fn cmsg_buf(&self) -> Buf {
self.cached_cmsg_bufs
.pop()
.unwrap_or_else(|| Buf::new(1024))
}
}
#[derive(Debug)]
struct IoUringTaskIds {
next: NumCell<u64>,
}
impl Default for IoUringTaskIds {
fn default() -> Self {
Self {
next: NumCell::new(1),
}
}
}
impl IoUringTaskIds {
fn next(&self) -> IoUringTaskId {
IoUringTaskId(self.next.fetch_add(1))
}
}
#[derive(Copy, Clone, Debug, Eq, PartialEq, Hash, Ord, PartialOrd)]
pub struct IoUringTaskId(u64);
impl IoUringTaskId {
#[allow(clippy::allow_attributes, dead_code)]
pub fn raw(&self) -> u64 {
self.0
}
#[allow(clippy::allow_attributes, dead_code)]
pub fn from_raw(id: u64) -> Self {
Self(id)
}
}
impl std::fmt::Display for IoUringTaskId {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Display::fmt(&self.0, f)
}
}
#[expect(clippy::derivable_impls)]
impl Default for IoUringTaskId {
fn default() -> Self {
Self(0)
}
}
struct Cancellable<'a> {
id: IoUringTaskId,
data: &'a IoUringData,
}
impl<'a> Drop for Cancellable<'a> {
fn drop(&mut self) {
self.data.cancel_task(self.id);
}
}

30
io-uring/src/ops.rs Normal file
View file

@ -0,0 +1,30 @@
use {crate::IoUringError, jay_utils::oserror::OsError};
pub mod accept;
pub mod async_cancel;
pub mod connect;
pub mod poll;
pub mod poll_external;
pub mod read_write;
pub mod read_write_no_cancel;
pub mod recvmsg;
pub mod sendmsg;
pub mod timeout;
pub mod timeout_external;
pub mod timeout_link;
pub type TaskResult<T> = Result<Result<T, OsError>, IoUringError>;
pub trait TaskResultExt<T> {
fn merge(self) -> Result<T, IoUringError>;
}
impl<T> TaskResultExt<T> for TaskResult<T> {
fn merge(self) -> Result<T, IoUringError> {
match self {
Ok(Ok(t)) => Ok(t),
Ok(Err(e)) => Err(IoUringError::OsError(e)),
Err(e) => Err(e),
}
}
}

67
io-uring/src/ops/accept.rs Normal file
View file

@ -0,0 +1,67 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, TaskResultExt,
pending_result::PendingResult,
sys::{IORING_OP_ACCEPT, io_uring_sqe},
},
std::rc::Rc,
uapi::{OwnedFd, c},
};
impl IoUring {
pub async fn accept(
&self,
fd: &Rc<OwnedFd>,
flags: c::c_int,
) -> Result<Rc<OwnedFd>, IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut pw = self.ring.cached_accepts.pop().unwrap_or_default();
pw.id = id.id;
pw.fd = fd.raw() as _;
pw.flags = flags as _;
pw.data = Some(Data {
pr: pr.clone(),
_fd: fd.clone(),
});
self.ring.schedule(pw);
}
Ok(pr.await.map(OwnedFd::new).map(Rc::new)).merge()
}
}
struct Data {
pr: PendingResult,
_fd: Rc<OwnedFd>,
}
#[derive(Default)]
pub struct AcceptTask {
id: IoUringTaskId,
fd: i32,
flags: u32,
data: Option<Data>,
}
unsafe impl Task for AcceptTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(data) = self.data.take() {
data.pr.complete(res);
}
ring.cached_accepts.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_ACCEPT;
sqe.fd = self.fd;
sqe.u2.addr = 0;
sqe.u1.addr2 = 0;
sqe.u3.accept_flags = self.flags;
}
}

48
io-uring/src/ops/async_cancel.rs Normal file
View file

@ -0,0 +1,48 @@
use {
crate::{
IoUringData, IoUringTaskId, Task,
sys::{IORING_OP_ASYNC_CANCEL, io_uring_sqe},
},
jay_utils::errorfmt::ErrorFmt,
uapi::c,
};
#[derive(Default)]
pub struct AsyncCancelTask {
id: IoUringTaskId,
target: IoUringTaskId,
}
impl IoUringData {
pub fn cancel_task_in_kernel(&self, target: IoUringTaskId) {
let id = self.id_raw();
let mut task = self.cached_cancels.pop().unwrap_or_default();
task.id = id;
task.target = target;
self.schedule(task);
}
}
unsafe impl Task for AsyncCancelTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(self: Box<Self>, ring: &IoUringData, res: i32) {
if let Err(e) = map_err!(res) {
if e.0 != c::ENOENT {
log::debug!("Could not cancel task: {}", ErrorFmt(e));
}
}
ring.cached_cancels.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_ASYNC_CANCEL;
sqe.u2.addr = self.target.raw();
}
fn is_cancel(&self) -> bool {
true
}
}

77
io-uring/src/ops/connect.rs Normal file
View file

@ -0,0 +1,77 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, TaskResultExt,
pending_result::PendingResult,
sys::{IORING_OP_CONNECT, io_uring_sqe},
},
std::{ptr, rc::Rc},
uapi::{OwnedFd, SockAddr, c},
};
impl IoUring {
pub async fn connect<T: SockAddr>(&self, fd: &Rc<OwnedFd>, t: &T) -> Result<(), IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut pw = self.ring.cached_connects.pop().unwrap_or_default();
pw.id = id.id;
pw.fd = fd.raw() as _;
unsafe {
ptr::copy_nonoverlapping(t, &mut pw.sockaddr as *mut _ as *mut _, 1);
}
pw.addrlen = size_of::<T>() as _;
pw.data = Some(Data {
pr: pr.clone(),
_fd: fd.clone(),
});
self.ring.schedule(pw);
}
Ok(pr.await.map(drop)).merge()
}
}
struct Data {
pr: PendingResult,
_fd: Rc<OwnedFd>,
}
pub struct ConnectTask {
id: IoUringTaskId,
fd: i32,
sockaddr: c::sockaddr_storage,
addrlen: u64,
data: Option<Data>,
}
impl Default for ConnectTask {
fn default() -> Self {
Self {
id: Default::default(),
fd: 0,
sockaddr: uapi::pod_zeroed(),
addrlen: 0,
data: None,
}
}
}
unsafe impl Task for ConnectTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(data) = self.data.take() {
data.pr.complete(res);
}
ring.cached_connects.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_CONNECT;
sqe.fd = self.fd;
sqe.u2.addr = &self.sockaddr as *const _ as _;
sqe.u1.off = self.addrlen;
}
}

70
io-uring/src/ops/poll.rs Normal file
View file

@ -0,0 +1,70 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, TaskResultExt,
ops::TaskResult,
pending_result::PendingResult,
sys::{IORING_OP_POLL_ADD, io_uring_sqe},
},
std::rc::Rc,
uapi::{OwnedFd, c},
};
impl IoUring {
pub async fn poll(&self, fd: &Rc<OwnedFd>, events: c::c_short) -> TaskResult<c::c_short> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut pw = self.ring.cached_polls.pop().unwrap_or_default();
pw.id = id.id;
pw.fd = fd.raw() as _;
pw.events = events as _;
pw.data = Some(Data {
pr: pr.clone(),
_fd: fd.clone(),
});
self.ring.schedule(pw);
}
Ok(pr.await.map(|v| v as c::c_short))
}
pub async fn readable(&self, fd: &Rc<OwnedFd>) -> Result<c::c_short, IoUringError> {
self.poll(fd, c::POLLIN).await.merge()
}
pub async fn writable(&self, fd: &Rc<OwnedFd>) -> Result<c::c_short, IoUringError> {
self.poll(fd, c::POLLOUT).await.merge()
}
}
struct Data {
pr: PendingResult,
_fd: Rc<OwnedFd>,
}
#[derive(Default)]
pub struct PollTask {
id: IoUringTaskId,
events: u16,
fd: i32,
data: Option<Data>,
}
unsafe impl Task for PollTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(data) = self.data.take() {
data.pr.complete(res);
}
ring.cached_polls.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_POLL_ADD;
sqe.fd = self.fd;
sqe.u3.poll_events = self.events;
}
}

113
io-uring/src/ops/poll_external.rs Normal file
View file

@ -0,0 +1,113 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task,
sys::{IORING_OP_POLL_ADD, io_uring_sqe},
},
jay_utils::oserror::OsError,
std::{cell::Cell, rc::Rc},
uapi::{OwnedFd, c},
};
pub trait PollCallback {
fn completed(self: Rc<Self>, res: Result<c::c_short, OsError>);
}
pub struct PendingPoll {
data: Rc<IoUringData>,
shared: Rc<PollExternalTaskShared>,
id: IoUringTaskId,
}
impl Drop for PendingPoll {
fn drop(&mut self) {
if self.shared.id.get() != self.id {
return;
}
self.shared.callback.take();
self.data.cancel_task(self.id);
}
}
impl IoUring {
pub fn poll_external(
&self,
fd: &Rc<OwnedFd>,
events: c::c_short,
callback: Rc<dyn PollCallback>,
) -> Result<PendingPoll, IoUringError> {
self.ring.check_destroyed()?;
let mut pw = self.ring.cached_polls_external.pop().unwrap_or_default();
pw.shared.id.set(self.ring.id_raw());
pw.shared.callback.set(Some(callback));
pw.fd = fd.raw() as _;
pw.events = events as _;
pw.data = Some(Data { _fd: fd.clone() });
let pending = PendingPoll {
data: self.ring.clone(),
shared: pw.shared.clone(),
id: pw.shared.id.get(),
};
self.ring.schedule(pw);
Ok(pending)
}
pub fn readable_external(
&self,
fd: &Rc<OwnedFd>,
callback: Rc<dyn PollCallback>,
) -> Result<PendingPoll, IoUringError> {
self.poll_external(fd, c::POLLIN, callback)
}
pub fn writable_external(
&self,
fd: &Rc<OwnedFd>,
callback: Rc<dyn PollCallback>,
) -> Result<PendingPoll, IoUringError> {
self.poll_external(fd, c::POLLOUT, callback)
}
}
struct Data {
_fd: Rc<OwnedFd>,
}
#[derive(Default)]
struct PollExternalTaskShared {
id: Cell<IoUringTaskId>,
callback: Cell<Option<Rc<dyn PollCallback>>>,
}
#[derive(Default)]
pub struct PollExternalTask {
shared: Rc<PollExternalTaskShared>,
events: u16,
fd: i32,
data: Option<Data>,
}
unsafe impl Task for PollExternalTask {
fn id(&self) -> IoUringTaskId {
self.shared.id.get()
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
self.data.take();
self.shared.id.set(Default::default());
if let Some(cb) = self.shared.callback.take() {
let res = if res < 0 {
Err(OsError::from(-res as c::c_int))
} else {
Ok(res as _)
};
cb.completed(res)
}
ring.cached_polls_external.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_POLL_ADD;
sqe.fd = self.fd;
sqe.u3.poll_events = self.events;
}
}

100
io-uring/src/ops/read_write.rs Normal file
View file

@ -0,0 +1,100 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, TaskResultExt,
pending_result::PendingResult,
sys::{IORING_OP_READ, IORING_OP_WRITE, io_uring_sqe},
},
jay_time::Time,
jay_utils::buf::Buf,
std::rc::Rc,
uapi::{OwnedFd, c},
};
impl IoUring {
pub async fn read(&self, fd: &Rc<OwnedFd>, buf: Buf) -> Result<usize, IoUringError> {
self.perform(fd, buf, None, IORING_OP_READ).await
}
pub async fn write(
&self,
fd: &Rc<OwnedFd>,
buf: Buf,
timeout: Option<Time>,
) -> Result<usize, IoUringError> {
self.perform(fd, buf, timeout, IORING_OP_WRITE).await
}
async fn perform(
&self,
fd: &Rc<OwnedFd>,
buf: Buf,
timeout: Option<Time>,
opcode: u8,
) -> Result<usize, IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut pw = self.ring.cached_read_writes.pop().unwrap_or_default();
pw.opcode = opcode;
pw.id = id.id;
pw.has_timeout = timeout.is_some();
pw.fd = fd.raw();
pw.buf = buf.as_ptr() as _;
pw.len = buf.len();
pw.data = Some(ReadWriteTaskData {
_fd: fd.clone(),
_buf: buf,
res: pr.clone(),
});
self.ring.schedule(pw);
if let Some(time) = timeout {
self.schedule_timeout_link(time);
}
}
Ok(pr.await.map(|v| v as usize)).merge()
}
}
struct ReadWriteTaskData {
_fd: Rc<OwnedFd>,
_buf: Buf,
res: PendingResult,
}
#[derive(Default)]
pub struct ReadWriteTask {
id: IoUringTaskId,
has_timeout: bool,
fd: c::c_int,
buf: usize,
len: usize,
data: Option<ReadWriteTaskData>,
opcode: u8,
}
unsafe impl Task for ReadWriteTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(data) = self.data.take() {
data.res.complete(res);
}
ring.cached_read_writes.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = self.opcode;
sqe.fd = self.fd as _;
sqe.u1.off = !0;
sqe.u2.addr = self.buf as _;
sqe.u3.rw_flags = 0;
sqe.len = self.len as _;
}
fn has_timeout(&self) -> bool {
self.has_timeout
}
}

135
io-uring/src/ops/read_write_no_cancel.rs Normal file
View file

@ -0,0 +1,135 @@
#[cfg(test)]
mod tests;
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, TaskResultExt,
pending_result::PendingResult,
sys::{IORING_OP_READ, IORING_OP_WRITE, io_uring_sqe},
},
jay_time::Time,
run_on_drop::on_drop,
uapi::{Fd, c},
};
impl IoUring {
pub async fn read_no_cancel(
&self,
fd: Fd,
offset: usize,
buf: &mut [u8],
cancel: impl FnOnce(IoUringTaskId),
) -> Result<usize, IoUringError> {
self.perform_no_cancel(
fd,
offset,
buf.as_mut_ptr(),
buf.len(),
None,
IORING_OP_READ,
cancel,
)
.await
}
pub async fn write_no_cancel(
&self,
fd: Fd,
offset: usize,
buf: &[u8],
timeout: Option<Time>,
cancel: impl FnOnce(IoUringTaskId),
) -> Result<usize, IoUringError> {
self.perform_no_cancel(
fd,
offset,
buf.as_ptr() as _,
buf.len(),
timeout,
IORING_OP_WRITE,
cancel,
)
.await
}
async fn perform_no_cancel(
&self,
fd: Fd,
offset: usize,
buf: *mut u8,
len: usize,
timeout: Option<Time>,
opcode: u8,
cancel: impl FnOnce(IoUringTaskId),
) -> Result<usize, IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut pw = self
.ring
.cached_read_writes_no_cancel
.pop()
.unwrap_or_default();
pw.opcode = opcode;
pw.id = id.id;
pw.has_timeout = timeout.is_some();
pw.fd = fd.raw();
pw.offset = offset;
pw.buf = buf as _;
pw.len = len;
pw.data = Some(ReadWriteTaskData { res: pr.clone() });
self.ring.schedule(pw);
if let Some(time) = timeout {
self.schedule_timeout_link(time);
}
}
let panic = on_drop(|| panic!("Operation cannot be cancelled from userspace"));
cancel(id.id);
let res = Ok(pr.await.map(|v| v as usize)).merge();
panic.forget();
res
}
}
struct ReadWriteTaskData {
res: PendingResult,
}
#[derive(Default)]
pub struct ReadWriteNoCancelTask {
id: IoUringTaskId,
has_timeout: bool,
fd: c::c_int,
offset: usize,
buf: usize,
len: usize,
data: Option<ReadWriteTaskData>,
opcode: u8,
}
unsafe impl Task for ReadWriteNoCancelTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(data) = self.data.take() {
data.res.complete(res);
}
ring.cached_read_writes_no_cancel.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = self.opcode;
sqe.fd = self.fd as _;
sqe.u1.off = self.offset as _;
sqe.u2.addr = self.buf as _;
sqe.u3.rw_flags = 0;
sqe.len = self.len as _;
}
fn has_timeout(&self) -> bool {
self.has_timeout
}
}

46
io-uring/src/ops/read_write_no_cancel/tests.rs Normal file
View file

@ -0,0 +1,46 @@
use {
crate::{IoUring, IoUringError},
jay_async_engine::AsyncEngine,
jay_utils::{oserror::OsError, queue::AsyncQueue},
std::rc::Rc,
uapi::c::ECANCELED,
};
fn cancel(timeout: bool) {
let eng = AsyncEngine::new();
let ring = IoUring::new(&eng, 32).unwrap();
let ring2 = ring.clone();
let ring3 = ring.clone();
let queue = Rc::new(AsyncQueue::new());
let queue2 = queue.clone();
let _fut1 = eng.spawn("", async move {
let (read, _write) = uapi::pipe().unwrap();
let mut buf = [10];
let res = ring
.read_no_cancel(read.borrow(), !0, &mut buf, |id| queue.push(id))
.await;
assert!(matches!(
res.unwrap_err(),
IoUringError::OsError(OsError(ECANCELED))
));
ring.stop();
});
let _fut2 = eng.spawn("", async move {
let id = queue2.pop().await;
if timeout {
ring2.timeout(1).await.unwrap();
}
ring2.cancel(id);
});
ring3.run().unwrap();
}
#[test]
fn cancel_in_kernel() {
cancel(true);
}
#[test]
fn cancel_in_userspace() {
cancel(true);
}

129
io-uring/src/ops/recvmsg.rs Normal file
View file

@ -0,0 +1,129 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task,
pending_result::PendingResult,
sys::{IORING_OP_RECVMSG, io_uring_sqe},
},
jay_utils::buf::Buf,
std::{cell::Cell, collections::VecDeque, mem::MaybeUninit, rc::Rc},
uapi::{OwnedFd, c},
};
impl IoUring {
pub async fn recvmsg(
&self,
fd: &Rc<OwnedFd>,
bufs: &mut [Buf],
fds: &mut VecDeque<Rc<OwnedFd>>,
) -> Result<usize, IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
let mut cmsg = self.ring.cmsg_buf();
let cmsg_len;
{
let mut rm = self.ring.cached_recvmsg.pop().unwrap_or_default();
rm.iovecs.clear();
for buf in bufs {
rm.bufs.push(buf.clone());
rm.iovecs.push(c::iovec {
iov_base: buf.as_ptr() as _,
iov_len: buf.len() as _,
});
}
rm.id = id.id;
rm.fd = fd.raw();
rm.msghdr.msg_control = cmsg.as_ptr() as _;
rm.msghdr.msg_controllen = cmsg.len() as _;
rm.msghdr.msg_iov = rm.iovecs.as_mut_ptr();
rm.msghdr.msg_iovlen = rm.iovecs.len() as _;
rm.data = Some(Data {
_cmsg: cmsg.clone(),
_fd: fd.clone(),
pr: pr.clone(),
});
cmsg_len = rm.cmsg_len.clone();
self.ring.schedule(rm);
}
macro_rules! return_cmsg {
() => {
self.ring.cached_cmsg_bufs.push(cmsg);
};
}
match pr.await {
Ok(n) => {
let mut cmsg_data = &cmsg[..cmsg_len.get()];
while cmsg_data.len() > 0 {
let (_, hdr, data) = match uapi::cmsg_read(&mut cmsg_data) {
Ok(m) => m,
Err(_) => {
return_cmsg!();
return Err(IoUringError::InvalidCmsgData);
}
};
if (hdr.cmsg_level, hdr.cmsg_type) == (c::SOL_SOCKET, c::SCM_RIGHTS) {
fds.extend(uapi::pod_iter(data).unwrap().map(Rc::new));
}
}
return_cmsg!();
Ok(n as _)
}
Err(e) => {
return_cmsg!();
Err(IoUringError::OsError(e))
}
}
}
}
struct Data {
_cmsg: Buf,
_fd: Rc<OwnedFd>,
pr: PendingResult,
}
pub struct RecvmsgTask {
id: IoUringTaskId,
fd: c::c_int,
bufs: Vec<Buf>,
iovecs: Vec<c::iovec>,
msghdr: c::msghdr,
cmsg_len: Rc<Cell<usize>>,
data: Option<Data>,
}
impl Default for RecvmsgTask {
fn default() -> Self {
RecvmsgTask {
id: Default::default(),
fd: 0,
bufs: vec![],
iovecs: vec![],
msghdr: unsafe { MaybeUninit::zeroed().assume_init() },
cmsg_len: Rc::new(Cell::new(0)),
data: None,
}
}
}
unsafe impl Task for RecvmsgTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
self.cmsg_len.set(self.msghdr.msg_controllen as _);
self.bufs.clear();
if let Some(data) = self.data.take() {
data.pr.complete(res);
}
ring.cached_recvmsg.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_RECVMSG;
sqe.fd = self.fd as _;
sqe.u2.addr = &self.msghdr as *const _ as _;
sqe.u3.msg_flags = c::MSG_CMSG_CLOEXEC as _;
}
}

139
io-uring/src/ops/sendmsg.rs Normal file
View file

@ -0,0 +1,139 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task,
pending_result::PendingResult,
sys::{IORING_OP_SENDMSG, io_uring_sqe},
},
jay_time::Time,
jay_utils::{buf::Buf, compat::IovLength, vec_ext::UninitVecExt},
std::{mem::MaybeUninit, ptr, rc::Rc},
uapi::{OwnedFd, c},
};
impl IoUring {
pub async fn sendmsg_one(
&self,
fd: &Rc<OwnedFd>,
buf: Buf,
fds: Vec<Rc<OwnedFd>>,
timeout: Option<Time>,
) -> Result<usize, IoUringError> {
self.sendmsg(fd, &mut [buf], fds, timeout).await
}
pub async fn sendmsg(
&self,
fd: &Rc<OwnedFd>,
bufs: &mut [Buf],
fds: Vec<Rc<OwnedFd>>,
timeout: Option<Time>,
) -> Result<usize, IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut st = self.ring.cached_sendmsg.pop().unwrap_or_default();
st.fds = fds;
if st.fds.len() > 0 {
let mut fd_ids = self.ring.fd_ids_scratch.borrow_mut();
fd_ids.clear();
fd_ids.extend(st.fds.iter().map(|f| f.raw()));
let space = uapi::cmsg_space(size_of_val(&fd_ids[..]));
st.cmsg.clear();
st.cmsg.reserve(space);
st.cmsg.set_len_safe(space);
let mut hdr: c::cmsghdr = uapi::pod_zeroed();
hdr.cmsg_level = c::SOL_SOCKET;
hdr.cmsg_type = c::SCM_RIGHTS;
uapi::cmsg_write(&mut &mut st.cmsg[..], hdr, &fd_ids[..]).unwrap();
st.msghdr.msg_control = st.cmsg.as_ptr() as _;
st.msghdr.msg_controllen = st.cmsg.len() as _;
} else {
st.msghdr.msg_control = ptr::null_mut();
st.msghdr.msg_controllen = 0;
}
st.id = id.id;
st.fd = fd.raw();
st.bufs.clear();
st.bufs.extend(bufs.iter_mut().map(|b| b.clone()));
st.iovecs.clear();
st.iovecs.extend(bufs.iter().map(|b| c::iovec {
iov_base: b.as_ptr() as _,
iov_len: b.len(),
}));
st.msghdr.msg_iov = st.iovecs.as_ptr() as _;
st.msghdr.msg_iovlen = st.iovecs.len() as IovLength;
st.data = Some(SendmsgTaskData {
_fd: fd.clone(),
res: pr.clone(),
});
st.has_timeout = timeout.is_some();
self.ring.schedule(st);
if let Some(timeout) = timeout {
self.schedule_timeout_link(timeout);
}
}
Ok(pr.await? as _)
}
}
struct SendmsgTaskData {
_fd: Rc<OwnedFd>,
res: PendingResult,
}
pub struct SendmsgTask {
id: IoUringTaskId,
iovecs: Vec<c::iovec>,
msghdr: c::msghdr,
bufs: Vec<Buf>,
fd: i32,
has_timeout: bool,
fds: Vec<Rc<OwnedFd>>,
cmsg: Vec<MaybeUninit<u8>>,
data: Option<SendmsgTaskData>,
}
impl Default for SendmsgTask {
fn default() -> Self {
unsafe {
SendmsgTask {
id: Default::default(),
iovecs: vec![],
msghdr: MaybeUninit::zeroed().assume_init(),
bufs: vec![],
fd: 0,
has_timeout: false,
fds: vec![],
cmsg: vec![],
data: None,
}
}
}
}
unsafe impl Task for SendmsgTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
self.fds.clear();
self.bufs.clear();
if let Some(data) = self.data.take() {
data.res.complete(res);
}
ring.cached_sendmsg.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_SENDMSG;
sqe.fd = self.fd;
sqe.u2.addr = &self.msghdr as *const _ as _;
sqe.u3.msg_flags = c::MSG_NOSIGNAL as _;
}
fn has_timeout(&self) -> bool {
self.has_timeout
}
}

63
io-uring/src/ops/timeout.rs Normal file
View file

@ -0,0 +1,63 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task,
pending_result::PendingResult,
sys::{IORING_OP_TIMEOUT, IORING_TIMEOUT_ABS, io_uring_sqe},
},
uapi::c,
};
#[repr(C)]
#[derive(Default)]
pub(super) struct timespec64 {
pub tv_sec: i64,
pub tv_nsec: c::c_long,
}
#[derive(Default)]
pub struct TimeoutTask {
id: IoUringTaskId,
timespec: timespec64,
pr: Option<PendingResult>,
}
impl IoUring {
pub async fn timeout(&self, timeout_nsec: u64) -> Result<(), IoUringError> {
self.ring.check_destroyed()?;
let id = self.ring.id();
let pr = self.ring.pending_results.acquire();
{
let mut pw = self.ring.cached_timeouts.pop().unwrap_or_default();
pw.id = id.id;
pw.timespec = timespec64 {
tv_sec: (timeout_nsec / 1_000_000_000) as _,
tv_nsec: (timeout_nsec % 1_000_000_000) as _,
};
pw.pr = Some(pr.clone());
self.ring.schedule(pw);
}
let _ = pr.await;
Ok(())
}
}
unsafe impl Task for TimeoutTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(mut self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(pr) = self.pr.take() {
pr.complete(res);
}
ring.cached_timeouts.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_TIMEOUT;
sqe.u2.addr = &self.timespec as *const _ as _;
sqe.len = 1;
sqe.u3.timeout_flags = IORING_TIMEOUT_ABS;
sqe.u1.off = 0;
}
}

94
io-uring/src/ops/timeout_external.rs Normal file
View file

@ -0,0 +1,94 @@
use {
crate::{
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, ops::timeout::timespec64,
sys::{IORING_OP_TIMEOUT, IORING_TIMEOUT_ABS, io_uring_sqe},
},
jay_utils::oserror::OsError,
std::{cell::Cell, rc::Rc},
uapi::c,
};
pub trait TimeoutCallback {
fn completed(self: Rc<Self>, res: Result<(), OsError>, data: u64);
}
pub struct PendingTimeout {
data: Rc<IoUringData>,
shared: Rc<TimeoutExternalTaskShared>,
id: IoUringTaskId,
}
impl Drop for PendingTimeout {
fn drop(&mut self) {
if self.shared.id.get() != self.id {
return;
}
self.shared.callback.take();
self.data.cancel_task(self.id);
}
}
#[derive(Default)]
struct TimeoutExternalTaskShared {
id: Cell<IoUringTaskId>,
callback: Cell<Option<Rc<dyn TimeoutCallback>>>,
}
#[derive(Default)]
pub struct TimeoutExternalTask {
timespec: timespec64,
shared: Rc<TimeoutExternalTaskShared>,
data: u64,
}
impl IoUring {
pub fn timeout_external(
&self,
timeout_nsec: u64,
callback: Rc<dyn TimeoutCallback>,
data: u64,
) -> Result<PendingTimeout, IoUringError> {
self.ring.check_destroyed()?;
let mut pw = self.ring.cached_timeouts_external.pop().unwrap_or_default();
pw.shared.id.set(self.ring.id_raw());
pw.shared.callback.set(Some(callback));
pw.timespec = timespec64 {
tv_sec: (timeout_nsec / 1_000_000_000) as _,
tv_nsec: (timeout_nsec % 1_000_000_000) as _,
};
pw.data = data;
let pending = PendingTimeout {
data: self.ring.clone(),
shared: pw.shared.clone(),
id: pw.shared.id.get(),
};
self.ring.schedule(pw);
Ok(pending)
}
}
unsafe impl Task for TimeoutExternalTask {
fn id(&self) -> IoUringTaskId {
self.shared.id.get()
}
fn complete(self: Box<Self>, ring: &IoUringData, res: i32) {
if let Some(pr) = self.shared.callback.take() {
let res = if res == -c::ETIME {
Ok(())
} else {
map_err!(res).map(drop)
};
pr.completed(res, self.data);
}
ring.cached_timeouts_external.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_TIMEOUT;
sqe.u2.addr = &self.timespec as *const _ as _;
sqe.len = 1;
sqe.u3.timeout_flags = IORING_TIMEOUT_ABS;
sqe.u1.off = 0;
}
}

41
io-uring/src/ops/timeout_link.rs Normal file
View file

@ -0,0 +1,41 @@
use crate::{
IoUring, IoUringData, IoUringTaskId, Task, ops::timeout::timespec64,
sys::{IORING_OP_LINK_TIMEOUT, IORING_TIMEOUT_ABS, io_uring_sqe},
};
use jay_time::Time;
#[derive(Default)]
pub struct TimeoutLinkTask {
id: IoUringTaskId,
timespec: timespec64,
}
impl IoUring {
pub(super) fn schedule_timeout_link(&self, timeout: Time) {
let id = self.ring.id_raw();
{
let mut to = self.ring.cached_timeout_links.pop().unwrap_or_default();
to.id = id;
to.timespec.tv_sec = timeout.0.tv_sec as _;
to.timespec.tv_nsec = timeout.0.tv_nsec as _;
self.ring.schedule(to);
}
}
}
unsafe impl Task for TimeoutLinkTask {
fn id(&self) -> IoUringTaskId {
self.id
}
fn complete(self: Box<Self>, ring: &IoUringData, _res: i32) {
ring.cached_timeout_links.push(self);
}
fn encode(&self, sqe: &mut io_uring_sqe) {
sqe.opcode = IORING_OP_LINK_TIMEOUT;
sqe.u2.addr = &self.timespec as *const _ as _;
sqe.len = 1;
sqe.u3.timeout_flags = IORING_TIMEOUT_ABS;
}
}

120
io-uring/src/pending_result.rs Normal file
View file

@ -0,0 +1,120 @@
use {
jay_utils::{numcell::NumCell, oserror::OsError, ptr_ext::PtrExt, stack::Stack},
std::{
cell::Cell,
future::Future,
pin::Pin,
rc::{Rc, Weak},
task::{Context, Poll, Waker},
},
uapi::c,
};
#[derive(Default)]
pub struct PendingResults {
data: Rc<PendingResultsData>,
}
impl PendingResults {
pub fn acquire(&self) -> PendingResult {
let pr = self.data.unused.pop().unwrap_or_else(|| {
Box::into_raw(Box::new(PendingResultData {
rc: NumCell::new(0),
base: Rc::downgrade(&self.data),
waker: Cell::new(None),
res: Cell::new(None),
}))
});
unsafe {
let prr = pr.deref();
debug_assert_eq!(prr.rc.get(), 0);
prr.rc.fetch_add(1);
PendingResult { pr }
}
}
}
#[derive(Default)]
struct PendingResultsData {
unused: Stack<*mut PendingResultData>,
}
impl Drop for PendingResultsData {
fn drop(&mut self) {
while let Some(pr) = self.unused.pop() {
unsafe {
drop(Box::from_raw(pr));
}
}
}
}
struct PendingResultData {
rc: NumCell<u32>,
base: Weak<PendingResultsData>,
waker: Cell<Option<Waker>>,
res: Cell<Option<i32>>,
}
pub struct PendingResult {
pr: *mut PendingResultData,
}
impl PendingResult {
pub fn complete(&self, res: i32) {
unsafe {
let pr = self.pr.deref();
pr.res.set(Some(res));
if let Some(waker) = pr.waker.take() {
waker.wake();
}
}
}
}
impl Drop for PendingResult {
fn drop(&mut self) {
{
let pr = unsafe { self.pr.deref() };
if pr.rc.fetch_sub(1) != 1 {
return;
}
if let Some(base) = pr.base.upgrade() {
pr.waker.set(None);
pr.res.set(None);
base.unused.push(self.pr);
return;
}
}
unsafe {
drop(Box::from_raw(self.pr));
}
}
}
impl Clone for PendingResult {
fn clone(&self) -> Self {
let pr = unsafe { self.pr.deref() };
pr.rc.fetch_add(1);
Self { pr: self.pr }
}
}
impl Future for PendingResult {
type Output = Result<i32, OsError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let pr = unsafe { self.pr.deref() };
if let Some(res) = pr.res.take() {
let res = if res < 0 {
Err(OsError::from(-res as c::c_int))
} else {
Ok(res)
};
Poll::Ready(res)
} else {
pr.waker.set(Some(cx.waker().clone()));
Poll::Pending
}
}
}

415
io-uring/src/sys.rs Normal file
View file

@ -0,0 +1,415 @@
#![allow(non_camel_case_types, dead_code)]
use {
jay_utils::oserror::OsError,
std::mem::MaybeUninit,
uapi::{OwnedFd, c},
};
#[repr(C)]
#[derive(Copy, Clone)]
pub struct io_uring_sqe {
pub opcode: u8,
pub flags: u8,
pub ioprio: u16,
pub fd: i32,
pub u1: io_uring_sqe_union1,
pub u2: io_uring_sqe_union2,
pub len: u32,
pub u3: io_uring_sqe_union3,
pub user_data: u64,
pub u4: io_uring_sqe_union4,
pub personality: u16,
pub u5: io_uring_sqe_union5,
pub __pad2: [u64; 2],
}
impl Default for io_uring_sqe {
fn default() -> Self {
unsafe { MaybeUninit::zeroed().assume_init() }
}
}
#[repr(C)]
#[derive(Copy, Clone)]
pub union io_uring_sqe_union1 {
pub off: u64,
pub addr2: u64,
}
#[repr(C)]
#[derive(Copy, Clone)]
pub union io_uring_sqe_union2 {
pub addr: u64,
pub splice_off_in: u64,
}
#[repr(C)]
#[derive(Copy, Clone)]
pub union io_uring_sqe_union3 {
pub rw_flags: c::c_int,
pub fsync_flags: u32,
pub poll_events: u16,
pub poll32_events: u32,
pub sync_range_flags: u32,
pub msg_flags: u32,
pub timeout_flags: u32,
pub accept_flags: u32,
pub cancel_flags: u32,
pub open_flags: u32,
pub statx_flags: u32,
pub fadvise_advice: u32,
pub splice_flags: u32,
pub rename_flags: u32,
pub unlink_flags: u32,
pub hardlink_flags: u32,
}
#[repr(C, packed)]
#[derive(Copy, Clone)]
pub union io_uring_sqe_union4 {
pub buf_index: u16,
pub buf_group: u16,
}
#[repr(C)]
#[derive(Copy, Clone)]
pub union io_uring_sqe_union5 {
pub splice_fd_in: i32,
pub file_index: u32,
}
pub const IOSQE_FIXED_FILE_BIT: u8 = 0;
pub const IOSQE_IO_DRAIN_BIT: u8 = 1;
pub const IOSQE_IO_LINK_BIT: u8 = 2;
pub const IOSQE_IO_HARDLINK_BIT: u8 = 3;
pub const IOSQE_ASYNC_BIT: u8 = 4;
pub const IOSQE_BUFFER_SELECT_BIT: u8 = 5;
pub const IOSQE_CQE_SKIP_SUCCESS_BIT: u8 = 6;
pub const IOSQE_FIXED_FILE: u8 = 1 << IOSQE_FIXED_FILE_BIT;
pub const IOSQE_IO_DRAIN: u8 = 1 << IOSQE_IO_DRAIN_BIT;
pub const IOSQE_IO_LINK: u8 = 1 << IOSQE_IO_LINK_BIT;
pub const IOSQE_IO_HARDLINK: u8 = 1 << IOSQE_IO_HARDLINK_BIT;
pub const IOSQE_ASYNC: u8 = 1 << IOSQE_ASYNC_BIT;
pub const IOSQE_BUFFER_SELECT: u8 = 1 << IOSQE_BUFFER_SELECT_BIT;
pub const IOSQE_CQE_SKIP_SUCCESS: u8 = 1 << IOSQE_CQE_SKIP_SUCCESS_BIT;
pub const IORING_SETUP_IOPOLL: u32 = 1 << 0;
pub const IORING_SETUP_SQPOLL: u32 = 1 << 1;
pub const IORING_SETUP_SQ_AFF: u32 = 1 << 2;
pub const IORING_SETUP_CQSIZE: u32 = 1 << 3;
pub const IORING_SETUP_CLAMP: u32 = 1 << 4;
pub const IORING_SETUP_ATTACH_WQ: u32 = 1 << 5;
pub const IORING_SETUP_R_DISABLED: u32 = 1 << 6;
pub const IORING_SETUP_SUBMIT_ALL: u32 = 1 << 7;
pub const IORING_SETUP_COOP_TASKRUN: u32 = 1 << 8;
pub const IORING_SETUP_TASKRUN_FLAG: u32 = 1 << 9;
pub const IORING_SETUP_SQE128: u32 = 1 << 10;
pub const IORING_SETUP_CQE32: u32 = 1 << 11;
pub const IORING_SETUP_SINGLE_ISSUER: u32 = 1 << 12;
pub const IORING_SETUP_DEFER_TASKRUN: u32 = 1 << 13;
pub const IORING_SETUP_NO_MMAP: u32 = 1 << 14;
pub const IORING_SETUP_REGISTERED_FD_ONLY: u32 = 1 << 15;
pub const IORING_SETUP_NO_SQARRAY: u32 = 1 << 16;
pub const IORING_SETUP_HYBRID_IOPOLL: u32 = 1 << 17;
pub const IORING_OP_NOP: u8 = 0;
pub const IORING_OP_READV: u8 = 1;
pub const IORING_OP_WRITEV: u8 = 2;
pub const IORING_OP_FSYNC: u8 = 3;
pub const IORING_OP_READ_FIXED: u8 = 4;
pub const IORING_OP_WRITE_FIXED: u8 = 5;
pub const IORING_OP_POLL_ADD: u8 = 6;
pub const IORING_OP_POLL_REMOVE: u8 = 7;
pub const IORING_OP_SYNC_FILE_RANGE: u8 = 8;
pub const IORING_OP_SENDMSG: u8 = 9;
pub const IORING_OP_RECVMSG: u8 = 10;
pub const IORING_OP_TIMEOUT: u8 = 11;
pub const IORING_OP_TIMEOUT_REMOVE: u8 = 12;
pub const IORING_OP_ACCEPT: u8 = 13;
pub const IORING_OP_ASYNC_CANCEL: u8 = 14;
pub const IORING_OP_LINK_TIMEOUT: u8 = 15;
pub const IORING_OP_CONNECT: u8 = 16;
pub const IORING_OP_FALLOCATE: u8 = 17;
pub const IORING_OP_OPENAT: u8 = 18;
pub const IORING_OP_CLOSE: u8 = 19;
pub const IORING_OP_FILES_UPDATE: u8 = 20;
pub const IORING_OP_STATX: u8 = 21;
pub const IORING_OP_READ: u8 = 22;
pub const IORING_OP_WRITE: u8 = 23;
pub const IORING_OP_FADVISE: u8 = 24;
pub const IORING_OP_MADVISE: u8 = 25;
pub const IORING_OP_SEND: u8 = 26;
pub const IORING_OP_RECV: u8 = 27;
pub const IORING_OP_OPENAT2: u8 = 28;
pub const IORING_OP_EPOLL_CTL: u8 = 29;
pub const IORING_OP_SPLICE: u8 = 30;
pub const IORING_OP_PROVIDE_BUFFERS: u8 = 31;
pub const IORING_OP_REMOVE_BUFFERS: u8 = 32;
pub const IORING_OP_TEE: u8 = 33;
pub const IORING_OP_SHUTDOWN: u8 = 34;
pub const IORING_OP_RENAMEAT: u8 = 35;
pub const IORING_OP_UNLINKAT: u8 = 36;
pub const IORING_OP_MKDIRAT: u8 = 37;
pub const IORING_OP_SYMLINKAT: u8 = 38;
pub const IORING_OP_LINKAT: u8 = 39;
pub const IORING_OP_LAST: u8 = 40;
pub const IORING_FSYNC_DATASYNC: u32 = 1 << 0;
pub const IORING_TIMEOUT_ABS: u32 = 1 << 0;
pub const IORING_TIMEOUT_UPDATE: u32 = 1 << 1;
pub const IORING_TIMEOUT_BOOTTIME: u32 = 1 << 2;
pub const IORING_TIMEOUT_REALTIME: u32 = 1 << 3;
pub const IORING_LINK_TIMEOUT_UPDATE: u32 = 1 << 4;
pub const IORING_TIMEOUT_ETIME_SUCCESS: u32 = 1 << 5;
pub const IORING_TIMEOUT_CLOCK_MASK: u32 = IORING_TIMEOUT_BOOTTIME | IORING_TIMEOUT_REALTIME;
pub const IORING_TIMEOUT_UPDATE_MASK: u32 = IORING_TIMEOUT_UPDATE | IORING_LINK_TIMEOUT_UPDATE;
pub const SPLICE_F_FD_IN_FIXED: u32 = 1 << 31;
pub const IORING_POLL_ADD_MULTI: u32 = 1 << 0;
pub const IORING_POLL_UPDATE_EVENTS: u32 = 1 << 1;
pub const IORING_POLL_UPDATE_USER_DATA: u32 = 1 << 2;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_cqe {
pub user_data: u64,
pub res: i32,
pub flags: u32,
}
pub const IORING_CQE_F_BUFFER: u32 = 1 << 0;
pub const IORING_CQE_F_MORE: u32 = 1 << 1;
pub const IORING_CQE_BUFFER_SHIFT: u32 = 16;
pub const IORING_OFF_SQ_RING: u64 = 0;
pub const IORING_OFF_CQ_RING: u64 = 0x8000000;
pub const IORING_OFF_SQES: u64 = 0x10000000;
#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct io_sqring_offsets {
pub head: u32,
pub tail: u32,
pub ring_mask: u32,
pub ring_entries: u32,
pub flags: u32,
pub dropped: u32,
pub array: u32,
pub resv1: u32,
pub resv2: u64,
}
#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct io_cqring_offsets {
pub head: u32,
pub tail: u32,
pub ring_mask: u32,
pub ring_entries: u32,
pub overflow: u32,
pub cqes: u32,
pub flags: u32,
pub resv1: u32,
pub resv2: u64,
}
#[repr(C)]
#[derive(Debug, Copy, Clone, Default)]
pub struct io_uring_params {
pub sq_entries: u32,
pub cq_entries: u32,
pub flags: u32,
pub sq_thread_cpu: u32,
pub sq_thread_idle: u32,
pub features: u32,
pub wq_fd: u32,
pub resv: [u32; 3],
pub sq_off: io_sqring_offsets,
pub cq_off: io_cqring_offsets,
}
pub const IORING_SQ_NEED_WAKEUP: u32 = 1 << 0;
pub const IORING_SQ_CQ_OVERFLOW: u32 = 1 << 1;
pub const IORING_CQ_EVENTFD_DISABLED: u32 = 1 << 0;
pub const IORING_ENTER_GETEVENTS: c::c_uint = 1 << 0;
pub const IORING_ENTER_SQ_WAKEUP: c::c_uint = 1 << 1;
pub const IORING_ENTER_SQ_WAIT: c::c_uint = 1 << 2;
pub const IORING_ENTER_EXT_ARG: c::c_uint = 1 << 3;
pub const IORING_FEAT_SINGLE_MMAP: u32 = 1 << 0;
pub const IORING_FEAT_NODROP: u32 = 1 << 1;
pub const IORING_FEAT_SUBMIT_STABLE: u32 = 1 << 2;
pub const IORING_FEAT_RW_CUR_POS: u32 = 1 << 3;
pub const IORING_FEAT_CUR_PERSONALITY: u32 = 1 << 4;
pub const IORING_FEAT_FAST_POLL: u32 = 1 << 5;
pub const IORING_FEAT_POLL_32BITS: u32 = 1 << 6;
pub const IORING_FEAT_SQPOLL_NONFIXED: u32 = 1 << 7;
pub const IORING_FEAT_EXT_ARG: u32 = 1 << 8;
pub const IORING_FEAT_NATIVE_WORKERS: u32 = 1 << 9;
pub const IORING_FEAT_RSRC_TAGS: u32 = 1 << 10;
pub const IORING_FEAT_CQE_SKIP: u32 = 1 << 11;
pub const IORING_REGISTER_BUFFERS: c::c_uint = 0;
pub const IORING_UNREGISTER_BUFFERS: c::c_uint = 1;
pub const IORING_REGISTER_FILES: c::c_uint = 2;
pub const IORING_UNREGISTER_FILES: c::c_uint = 3;
pub const IORING_REGISTER_EVENTFD: c::c_uint = 4;
pub const IORING_UNREGISTER_EVENTFD: c::c_uint = 5;
pub const IORING_REGISTER_FILES_UPDATE: c::c_uint = 6;
pub const IORING_REGISTER_EVENTFD_ASYNC: c::c_uint = 7;
pub const IORING_REGISTER_PROBE: c::c_uint = 8;
pub const IORING_REGISTER_PERSONALITY: c::c_uint = 9;
pub const IORING_UNREGISTER_PERSONALITY: c::c_uint = 10;
pub const IORING_REGISTER_RESTRICTIONS: c::c_uint = 11;
pub const IORING_REGISTER_ENABLE_RINGS: c::c_uint = 12;
pub const IORING_REGISTER_FILES2: c::c_uint = 13;
pub const IORING_REGISTER_FILES_UPDATE2: c::c_uint = 14;
pub const IORING_REGISTER_BUFFERS2: c::c_uint = 15;
pub const IORING_REGISTER_BUFFERS_UPDATE: c::c_uint = 16;
pub const IORING_REGISTER_IOWQ_AFF: c::c_uint = 17;
pub const IORING_UNREGISTER_IOWQ_AFF: c::c_uint = 18;
pub const IORING_REGISTER_IOWQ_MAX_WORKERS: c::c_uint = 19;
pub const IO_WQ_BOUND: u32 = 0;
pub const IO_WQ_UNBOUND: u32 = 1;
#[repr(C, align(8))]
#[derive(Debug, Copy, Clone)]
pub struct AlignedU64(pub u64);
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_files_update {
pub offset: u32,
pub resv: u32,
pub fds: AlignedU64,
}
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_rsrc_register {
pub nr: u32,
pub resv: u32,
pub resv2: u64,
pub data: AlignedU64,
pub tags: AlignedU64,
}
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_rsrc_update {
pub offset: u32,
pub resv: u32,
pub data: AlignedU64,
}
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_rsrc_update2 {
pub offset: u32,
pub resv: u32,
pub data: AlignedU64,
pub tags: AlignedU64,
pub nr: u32,
pub resv2: u32,
}
pub const IORING_REGISTER_FILES_SKIP: i32 = -2;
pub const IO_URING_OP_SUPPORTED: u32 = 1 << 0;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_probe_op {
pub op: u8,
pub resv: u8,
pub flags: u16,
pub resv2: u32,
}
#[repr(C)]
#[derive(Debug)]
pub struct io_uring_probe {
pub last_op: u8,
pub ops_len: u8,
pub resv: u16,
pub resv2: [u32; 3],
pub ops: [io_uring_probe_op; 0],
}
#[repr(C)]
#[derive(Copy, Clone)]
pub struct io_uring_restriction {
pub opcode: u16,
pub u1: io_uring_restriction_union1,
pub resv: u8,
pub resv2: [u32; 3],
}
#[repr(C)]
#[derive(Copy, Clone)]
pub union io_uring_restriction_union1 {
pub register_op: u8,
pub sqe_op: u8,
pub sqe_flags: u8,
}
pub const IORING_RESTRICTION_REGISTER_OP: u16 = 0;
pub const IORING_RESTRICTION_SQE_OP: u16 = 1;
pub const IORING_RESTRICTION_SQE_FLAGS_ALLOWED: u16 = 2;
pub const IORING_RESTRICTION_SQE_FLAGS_REQUIRED: u16 = 3;
#[repr(C)]
#[derive(Debug, Copy, Clone)]
pub struct io_uring_getevents_arg {
sigmask: u64,
sigmask_sz: u32,
pad: u32,
ts: u64,
}
pub fn io_uring_setup(entries: u32, params: &mut io_uring_params) -> Result<OwnedFd, OsError> {
let res = unsafe {
c::syscall(
c::SYS_io_uring_setup,
entries as usize,
params as *mut _ as usize,
)
};
if res < 0 {
Err(OsError::default())
} else {
Ok(OwnedFd::new(res as _))
}
}
pub fn io_uring_enter(
fd: c::c_int,
to_submit: c::c_uint,
min_complete: c::c_uint,
flags: c::c_uint,
) -> Result<usize, OsError> {
let res = unsafe {
c::syscall(
c::SYS_io_uring_enter,
fd as usize,
to_submit as usize,
min_complete as usize,
flags as usize,
0usize,
0usize,
)
};
if res < 0 {
Err(OsError::default())
} else {
Ok(res as usize)
}
}