1
0
Fork 0
forked from wry/wry

io_uring: add ops for non-owning read/write operations

This commit is contained in:
Julian Orth 2024-09-07 17:11:19 +02:00
parent 5b36980e72
commit 874d0d0c59
4 changed files with 194 additions and 1 deletions

View file

@ -5,7 +5,8 @@ use {
io_uring::{
ops::{
accept::AcceptTask, async_cancel::AsyncCancelTask, connect::ConnectTask,
poll::PollTask, read_write::ReadWriteTask, recvmsg::RecvmsgTask,
poll::PollTask, read_write::ReadWriteTask,
read_write_no_cancel::ReadWriteNoCancelTask, recvmsg::RecvmsgTask,
sendmsg::SendmsgTask, timeout::TimeoutTask, timeout_link::TimeoutLinkTask,
},
pending_result::PendingResults,
@ -205,6 +206,7 @@ impl IoUring {
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_sendmsg: Default::default(),
@ -266,6 +268,7 @@ struct IoUringData {
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_sendmsg: Stack<Box<SendmsgTask>>,

View file

@ -5,6 +5,7 @@ pub mod async_cancel;
pub mod connect;
pub mod poll;
pub mod read_write;
pub mod read_write_no_cancel;
pub mod recvmsg;
pub mod sendmsg;
pub mod timeout;

View file

@ -0,0 +1,139 @@
#[cfg(test)]
mod tests;
use {
crate::{
io_uring::{
pending_result::PendingResult,
sys::{io_uring_sqe, IORING_OP_READ, IORING_OP_WRITE},
IoUring, IoUringData, IoUringError, IoUringTaskId, Task, TaskResultExt,
},
time::Time,
utils::on_drop::OnDrop,
},
uapi::{c, Fd},
};
impl IoUring {
#[expect(dead_code)]
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
}
#[expect(dead_code)]
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 = OnDrop(|| 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
}
}

View file

@ -0,0 +1,50 @@
use {
crate::{
async_engine::AsyncEngine,
io_uring::{IoUring, IoUringError},
utils::{oserror::OsError, queue::AsyncQueue},
wheel::Wheel,
},
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 wheel = Wheel::new(&eng, &ring).unwrap();
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 {
wheel.timeout(1).await.unwrap();
}
ring2.cancel(id);
});
ring3.run().unwrap();
}
#[test]
fn cancel_in_kernel() {
cancel(true);
}
#[test]
fn cancel_in_userspace() {
cancel(true);
}