#[cfg_attr(all(doc, docsrs), doc(cfg(all())))]
#[allow(unused_imports)]
pub use std::os::fd::{AsRawFd, FromRawFd, IntoRawFd, RawFd};
use std::{
collections::VecDeque, io, os::fd::OwnedFd, pin::Pin, ptr::NonNull, sync::Arc, task::Poll,
time::Duration,
};
use compio_log::{instrument, trace};
use crossbeam_queue::SegQueue;
cfg_if::cfg_if! {
if #[cfg(feature = "io-uring-cqe32")] {
use io_uring::cqueue::Entry32 as CEntry;
} else {
use io_uring::cqueue::Entry as CEntry;
}
}
cfg_if::cfg_if! {
if #[cfg(feature = "io-uring-sqe128")] {
use io_uring::squeue::Entry128 as SEntry;
} else {
use io_uring::squeue::Entry as SEntry;
}
}
use io_uring::{
opcode::{AsyncCancel, PollAdd},
types::{Fd, SubmitArgs, Timespec},
IoUring,
};
pub(crate) use libc::{sockaddr_storage, socklen_t};
use slab::Slab;
use crate::{syscall, AsyncifyPool, Entry, OutEntries, ProactorBuilder};
pub(crate) mod op;
pub(crate) use crate::unix::RawOp;
pub enum OpEntry {
Submission(io_uring::squeue::Entry),
#[cfg(feature = "io-uring-sqe128")]
Submission128(io_uring::squeue::Entry128),
Blocking,
}
impl From<io_uring::squeue::Entry> for OpEntry {
fn from(value: io_uring::squeue::Entry) -> Self {
Self::Submission(value)
}
}
#[cfg(feature = "io-uring-sqe128")]
impl From<io_uring::squeue::Entry128> for OpEntry {
fn from(value: io_uring::squeue::Entry128) -> Self {
Self::Submission128(value)
}
}
pub trait OpCode {
fn create_entry(self: Pin<&mut Self>) -> OpEntry;
fn call_blocking(self: Pin<&mut Self>) -> io::Result<usize> {
unreachable!("this operation is asynchronous")
}
}
pub(crate) struct Driver {
inner: IoUring<SEntry, CEntry>,
squeue: VecDeque<SEntry>,
notifier: Notifier,
pool: AsyncifyPool,
pool_completed: Arc<SegQueue<Entry>>,
}
impl Driver {
const CANCEL: u64 = u64::MAX;
const NOTIFY: u64 = u64::MAX - 1;
pub fn new(builder: &ProactorBuilder) -> io::Result<Self> {
instrument!(compio_log::Level::TRACE, "new", ?builder);
trace!("new iour driver");
let mut squeue = VecDeque::with_capacity(builder.capacity as usize);
let notifier = Notifier::new()?;
#[allow(clippy::useless_conversion)]
squeue.push_back(
PollAdd::new(Fd(notifier.as_raw_fd()), libc::POLLIN as _)
.multi(true)
.build()
.user_data(Self::NOTIFY)
.into(),
);
Ok(Self {
inner: IoUring::builder().build(builder.capacity)?,
squeue,
notifier,
pool: builder.create_or_get_thread_pool(),
pool_completed: Arc::new(SegQueue::new()),
})
}
fn submit_auto(&mut self, timeout: Option<Duration>, wait: bool) -> io::Result<()> {
instrument!(compio_log::Level::TRACE, "submit_auto", ?timeout, wait);
let res = if wait {
if let Some(duration) = timeout {
let timespec = timespec(duration);
let args = SubmitArgs::new().timespec(×pec);
self.inner.submitter().submit_with_args(1, &args)
} else {
self.inner.submit_and_wait(1)
}
} else {
self.inner.submit()
};
trace!("submit result: {res:?}");
match res {
Ok(_) => {
if self.inner.completion().is_empty() {
Err(io::Error::from_raw_os_error(libc::ETIMEDOUT))
} else {
Ok(())
}
}
Err(e) => match e.raw_os_error() {
Some(libc::ETIME) => Err(io::Error::from_raw_os_error(libc::ETIMEDOUT)),
Some(libc::EBUSY) | Some(libc::EAGAIN) => Ok(()),
_ => Err(e),
},
}
}
fn flush_submissions(&mut self) -> bool {
instrument!(compio_log::Level::TRACE, "flush_submissions");
let mut ended_ops = false;
let mut inner_squeue = self.inner.submission();
while !inner_squeue.is_full() {
if self.squeue.len() <= inner_squeue.capacity() - inner_squeue.len() {
trace!("inner_squeue have enough space, flush all entries");
let (s1, s2) = self.squeue.as_slices();
unsafe {
inner_squeue
.push_multiple(s1)
.expect("queue has enough space");
inner_squeue
.push_multiple(s2)
.expect("queue has enough space");
}
self.squeue.clear();
ended_ops = true;
break;
} else if let Some(entry) = self.squeue.pop_front() {
trace!("inner_squeue have not enough space, flush an entry");
unsafe { inner_squeue.push(&entry) }.expect("queue has enough space");
} else {
trace!("self.squeue is empty, skip");
ended_ops = true;
break;
}
}
inner_squeue.sync();
ended_ops
}
fn poll_entries(&mut self, entries: &mut impl Extend<Entry>) {
while let Some(entry) = self.pool_completed.pop() {
entries.extend(Some(entry));
}
let mut cqueue = self.inner.completion();
cqueue.sync();
let completed_entries = cqueue.filter_map(|entry| match entry.user_data() {
Self::CANCEL => None,
Self::NOTIFY => {
const IORING_CQE_F_MORE: u32 = 1 << 1;
let flags = entry.flags();
debug_assert!(flags & IORING_CQE_F_MORE == IORING_CQE_F_MORE);
self.notifier.clear().expect("cannot clear notifier");
None
}
_ => Some(create_entry(entry)),
});
entries.extend(completed_entries);
}
pub fn create_op<T: crate::sys::OpCode + 'static>(&self, user_data: usize, op: T) -> RawOp {
RawOp::new(user_data, op)
}
pub fn attach(&mut self, _fd: RawFd) -> io::Result<()> {
Ok(())
}
pub fn cancel(&mut self, user_data: usize, _registry: &mut Slab<RawOp>) {
instrument!(compio_log::Level::TRACE, "cancel", user_data);
trace!("cancel RawOp");
#[allow(clippy::useless_conversion)]
self.squeue.push_back(
AsyncCancel::new(user_data as _)
.build()
.user_data(Self::CANCEL)
.into(),
);
}
pub fn push(&mut self, user_data: usize, op: &mut RawOp) -> Poll<io::Result<usize>> {
instrument!(compio_log::Level::TRACE, "push", user_data);
let op_pin = op.as_pin();
trace!("push RawOp");
match op_pin.create_entry() {
OpEntry::Submission(entry) => {
#[allow(clippy::useless_conversion)]
self.squeue
.push_back(entry.user_data(user_data as _).into());
Poll::Pending
}
#[cfg(feature = "io-uring-sqe128")]
OpEntry::Submission128(_entry) => {
self.squeue.push_back(_entry.user_data(user_data as _));
Poll::Pending
}
OpEntry::Blocking => {
if self.push_blocking(user_data, op)? {
Poll::Pending
} else {
Poll::Ready(Err(io::Error::from_raw_os_error(libc::EBUSY)))
}
}
}
}
fn push_blocking(&mut self, user_data: usize, op: &mut RawOp) -> io::Result<bool> {
struct SendWrapper<T>(T);
unsafe impl<T> Send for SendWrapper<T> {}
let op = SendWrapper(NonNull::from(op));
let handle = self.handle()?;
let completed = self.pool_completed.clone();
let is_ok = self
.pool
.dispatch(move || {
#[allow(clippy::redundant_locals)]
let mut op = op;
let op = unsafe { op.0.as_mut() };
let op_pin = op.as_pin();
let res = op_pin.call_blocking();
completed.push(Entry::new(user_data, res));
handle.notify().ok();
})
.is_ok();
Ok(is_ok)
}
pub unsafe fn poll(
&mut self,
timeout: Option<Duration>,
mut entries: OutEntries<impl Extend<usize>>,
) -> io::Result<()> {
instrument!(compio_log::Level::TRACE, "poll", ?timeout);
trace!("start polling");
loop {
let ended = self.flush_submissions();
self.submit_auto(timeout, ended)?;
self.poll_entries(&mut entries);
if ended {
trace!("polling ended");
break;
}
}
Ok(())
}
pub fn handle(&self) -> io::Result<NotifyHandle> {
self.notifier.handle()
}
}
impl AsRawFd for Driver {
fn as_raw_fd(&self) -> RawFd {
self.inner.as_raw_fd()
}
}
fn create_entry(entry: CEntry) -> Entry {
let result = entry.result();
let result = if result < 0 {
let result = if result == -libc::ECANCELED {
libc::ETIMEDOUT
} else {
-result
};
Err(io::Error::from_raw_os_error(result))
} else {
Ok(result as _)
};
Entry::new(entry.user_data() as _, result)
}
fn timespec(duration: std::time::Duration) -> Timespec {
Timespec::new()
.sec(duration.as_secs())
.nsec(duration.subsec_nanos())
}
#[derive(Debug)]
struct Notifier {
fd: OwnedFd,
}
impl Notifier {
fn new() -> io::Result<Self> {
let fd = syscall!(libc::eventfd(0, libc::EFD_CLOEXEC | libc::EFD_NONBLOCK))?;
let fd = unsafe { OwnedFd::from_raw_fd(fd) };
Ok(Self { fd })
}
pub fn clear(&self) -> io::Result<()> {
loop {
let mut buffer = [0u64];
let res = syscall!(libc::read(
self.fd.as_raw_fd(),
buffer.as_mut_ptr().cast(),
std::mem::size_of::<u64>()
));
match res {
Ok(len) => {
debug_assert_eq!(len, std::mem::size_of::<u64>() as _);
break Ok(());
}
Err(e) if e.kind() == io::ErrorKind::WouldBlock => break Ok(()),
Err(e) if e.kind() == io::ErrorKind::Interrupted => continue,
Err(e) => break Err(e),
}
}
}
pub fn handle(&self) -> io::Result<NotifyHandle> {
Ok(NotifyHandle::new(self.fd.try_clone()?))
}
}
impl AsRawFd for Notifier {
fn as_raw_fd(&self) -> RawFd {
self.fd.as_raw_fd()
}
}
pub struct NotifyHandle {
fd: OwnedFd,
}
impl NotifyHandle {
pub(crate) fn new(fd: OwnedFd) -> Self {
Self { fd }
}
pub fn notify(&self) -> io::Result<()> {
let data = 1u64;
syscall!(libc::write(
self.fd.as_raw_fd(),
&data as *const _ as *const _,
std::mem::size_of::<u64>(),
))?;
Ok(())
}
}