use std::collections::{HashMap, HashSet, VecDeque};
use std::fs::File;
use std::io;
use std::io::Write as _;
use std::os::fd::{AsRawFd, FromRawFd};
use std::os::unix::ffi::OsStrExt;
use std::pin::Pin;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::mpsc::{self, TryRecvError};
use std::task::{Context, Poll};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use io_uring::{IoUring, opcode, types};
const DRAIN_TIMEOUT: Duration = Duration::from_secs(5);
use tokio::sync::{OwnedSemaphorePermit, Semaphore, TryAcquireError, oneshot};
const CANCEL_BIT: u64 = 1 << 63;
const CURRENT_POSITION: u64 = u64::MAX;
const MAX_READ_LEN: usize = 0x7fff_f000;
fn aligned_geometry(offset: u64, len: usize, align: usize) -> Option<(u64, usize, usize)> {
if align == 0 || !align.is_power_of_two() || align > MAX_READ_LEN {
return None;
}
if align == 1 {
return Some((offset, 0, len));
}
let mask = align as u64 - 1;
let kernel_offset = offset & !mask;
let head = usize::try_from(offset - kernel_offset).ok()?;
let region_len = head.checked_add(len)?.checked_next_multiple_of(align)?;
Some((kernel_offset, head, region_len))
}
const LOOP_HEARTBEAT: Duration = Duration::from_millis(50);
const IDLE_HEARTBEAT: Duration = Duration::from_secs(1);
const IOWQ_MAX_BOUNDED_WORKERS: u32 = 16;
const MAX_CONSECUTIVE_SUBMIT_ERRORS: u32 = 128;
const MAX_TRANSIENT_RETRIES: u32 = 16;
struct EventFd {
fd: std::os::fd::RawFd,
}
impl EventFd {
fn new() -> io::Result<Self> {
let fd = unsafe { libc::eventfd(0, libc::EFD_NONBLOCK | libc::EFD_CLOEXEC) };
if fd < 0 {
return Err(io::Error::last_os_error());
}
Ok(Self { fd })
}
fn as_raw(&self) -> std::os::fd::RawFd {
self.fd
}
fn signal(&self) {
let v: u64 = 1;
unsafe {
libc::write(self.fd, (&v as *const u64).cast(), 8);
}
}
fn drain(&self) {
let mut v: u64 = 0;
while unsafe { libc::read(self.fd, (&mut v as *mut u64).cast(), 8) } == 8 {}
}
}
impl Drop for EventFd {
fn drop(&mut self) {
unsafe {
libc::close(self.fd);
}
}
}
fn wait_for_events(cq: &EventFd, wake: &EventFd, timeout: Duration) {
let mut fds = [
libc::pollfd {
fd: cq.as_raw(),
events: libc::POLLIN,
revents: 0,
},
libc::pollfd {
fd: wake.as_raw(),
events: libc::POLLIN,
revents: 0,
},
];
let ms = timeout.as_millis().min(i32::MAX as u128) as libc::c_int;
unsafe {
libc::poll(fds.as_mut_ptr(), fds.len() as libc::nfds_t, ms);
}
}
#[derive(Debug)]
pub enum ProbeFailure {
Setup(io::Error),
ReadOp(io::Error),
}
impl ProbeFailure {
pub fn is_expected_restriction(&self) -> bool {
let err = match self {
ProbeFailure::Setup(e) | ProbeFailure::ReadOp(e) => e,
};
matches!(
err.raw_os_error(),
Some(libc::EACCES) | Some(libc::EPERM) | Some(libc::ENOSYS) | Some(libc::EINVAL) | Some(libc::EOPNOTSUPP)
)
}
}
type AcquireFut = Pin<Box<dyn Future<Output = Result<OwnedSemaphorePermit, tokio::sync::AcquireError>> + Send>>;
#[derive(Default)]
struct DriverStats {
submitted: AtomicU64,
delivered: AtomicU64,
orphan_reclaimed: AtomicU64,
in_flight: AtomicU64,
cancel_succeeded: AtomicU64,
cancel_not_found: AtomicU64,
cancel_already: AtomicU64,
cq_overflow: AtomicU64,
submit_errors: AtomicU64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub struct StatsSnapshot {
pub submitted: u64,
pub delivered: u64,
pub orphan_reclaimed: u64,
pub in_flight: u64,
pub cancel_succeeded: u64,
pub cancel_not_found: u64,
pub cancel_already: u64,
pub cq_overflow: u64,
pub submit_errors: u64,
}
enum Msg {
Read {
id: u64,
file: Arc<File>,
offset: u64,
len: usize,
done: oneshot::Sender<io::Result<Vec<u8>>>,
permit: OwnedSemaphorePermit,
align: usize,
},
Cancel {
id: u64,
},
Shutdown,
#[cfg(feature = "fault-injection")]
TestPanic,
}
struct Pending {
buf: Vec<u8>,
file: Arc<File>,
done: Option<oneshot::Sender<io::Result<Vec<u8>>>>,
offset: u64,
nread: usize,
_permit: OwnedSemaphorePermit,
pad: usize,
head: usize,
want: usize,
region_len: usize,
align: usize,
transient_retries: u32,
}
impl Pending {
fn read_sqe(&self, ud: u64) -> io_uring::squeue::Entry {
let remaining = self.region_len - self.nread;
let ptr = unsafe { self.buf.as_ptr().add(self.pad + self.nread).cast_mut() };
let next_off = self.offset + self.nread as u64;
opcode::Read::new(types::Fd(self.file.as_raw_fd()), ptr, remaining as u32)
.offset(next_off)
.build()
.user_data(ud)
}
}
enum HandleState {
Inert,
WaitingPermit {
acquire: AcquireFut,
file: Arc<File>,
offset: u64,
len: usize,
align: usize,
done: oneshot::Sender<io::Result<Vec<u8>>>,
wake: Arc<EventFd>,
},
Submitted {
wake: Arc<EventFd>,
},
}
pub struct ReadHandle {
id: u64,
rx: oneshot::Receiver<io::Result<Vec<u8>>>,
tx: mpsc::Sender<Msg>,
finished: bool,
cancel_on_drop: bool,
state: HandleState,
}
impl ReadHandle {
pub fn without_cancel_on_drop(mut self) -> Self {
self.cancel_on_drop = false;
self
}
}
impl Future for ReadHandle {
type Output = io::Result<Vec<u8>>;
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = &mut *self;
let acquired = match &mut this.state {
HandleState::WaitingPermit { acquire, .. } => match acquire.as_mut().poll(cx) {
Poll::Pending => return Poll::Pending,
Poll::Ready(res) => Some(res),
},
_ => None,
};
if let Some(res) = acquired {
let Ok(permit) = res else {
this.finished = true;
return Poll::Ready(Err(io::Error::other("uring driver shut down")));
};
let submitted_wake = match &this.state {
HandleState::WaitingPermit { wake, .. } => Arc::clone(wake),
_ => unreachable!("state was WaitingPermit"),
};
let HandleState::WaitingPermit {
file,
offset,
len,
align,
done,
wake,
..
} = std::mem::replace(&mut this.state, HandleState::Submitted { wake: submitted_wake })
else {
unreachable!("state was WaitingPermit")
};
if this
.tx
.send(Msg::Read {
id: this.id,
file,
offset,
len,
done,
permit,
align,
})
.is_err()
{
this.finished = true;
return Poll::Ready(Err(io::Error::other("uring driver shut down")));
}
wake.signal();
}
match Pin::new(&mut this.rx).poll(cx) {
Poll::Ready(res) => {
this.finished = true;
Poll::Ready(match res {
Ok(inner) => inner,
Err(_) => Err(io::Error::other("uring driver shut down before completion")),
})
}
Poll::Pending => Poll::Pending,
}
}
}
impl Drop for ReadHandle {
fn drop(&mut self) {
if let HandleState::Submitted { wake } = &self.state
&& !self.finished
&& self.cancel_on_drop
{
self.rx.close();
let _ = self.tx.send(Msg::Cancel { id: self.id });
wake.signal();
}
}
}
struct Shard {
tx: mpsc::Sender<Msg>,
handle: Option<JoinHandle<()>>,
stats: Arc<DriverStats>,
sem: Arc<Semaphore>,
wake_efd: Arc<EventFd>,
}
impl Shard {
fn join(&mut self) {
if let Some(h) = self.handle.take() {
let _ = self.tx.send(Msg::Shutdown);
self.wake_efd.signal();
let _ = h.join();
}
}
}
impl Drop for Shard {
fn drop(&mut self) {
self.join();
}
}
pub struct UringDriver {
shards: Vec<Shard>,
next_id: AtomicU64,
rr: AtomicUsize,
}
impl UringDriver {
pub fn probe_and_start(entries: u32) -> Result<Self, ProbeFailure> {
Self::probe_and_start_sharded(entries, 1)
}
pub fn probe_and_start_sharded(entries: u32, shards: usize) -> Result<Self, ProbeFailure> {
let mut started = Vec::with_capacity(shards.max(1));
for i in 0..shards.max(1) {
started.push(Self::start_shard(entries, i == 0)?);
}
Ok(Self {
shards: started,
next_id: AtomicU64::new(1),
rr: AtomicUsize::new(0),
})
}
fn shard(&self) -> &Shard {
let n = self.shards.len();
&self.shards[self.rr.fetch_add(1, Ordering::Relaxed) % n]
}
fn start_shard(entries: u32, probe: bool) -> Result<Shard, ProbeFailure> {
let mut ring = IoUring::new(entries).map_err(ProbeFailure::Setup)?;
if !ring.params().is_feature_nodrop() {
return Err(ProbeFailure::Setup(io::Error::from_raw_os_error(libc::ENOSYS)));
}
if probe {
probe_real_read(&mut ring).map_err(ProbeFailure::ReadOp)?;
}
let cq_efd = EventFd::new().map_err(ProbeFailure::Setup)?;
ring.submitter()
.register_eventfd(cq_efd.as_raw())
.map_err(ProbeFailure::Setup)?;
let mut iowq_max = [IOWQ_MAX_BOUNDED_WORKERS, 0u32];
let _ = ring.submitter().register_iowq_max_workers(&mut iowq_max);
let wake_efd = Arc::new(EventFd::new().map_err(ProbeFailure::Setup)?);
let thread_wake = Arc::clone(&wake_efd);
let (tx, rx) = mpsc::channel();
let stats = Arc::new(DriverStats::default());
let thread_stats = Arc::clone(&stats);
let sem = Arc::new(Semaphore::new(entries as usize));
let thread_sem = Arc::clone(&sem);
#[cfg(feature = "fault-injection")]
if std::env::var_os("RUSTFS_URING_FAULT_SPAWN").is_some() {
return Err(ProbeFailure::Setup(io::Error::from_raw_os_error(libc::EAGAIN)));
}
let handle = std::thread::Builder::new()
.name("uring-spike-driver".into())
.spawn(move || drive(ring, rx, thread_stats, thread_sem, cq_efd, thread_wake))
.map_err(ProbeFailure::Setup)?;
Ok(Shard {
tx,
handle: Some(handle),
stats,
sem,
wake_efd,
})
}
pub fn read_at(&self, file: Arc<File>, offset: u64, len: usize) -> ReadHandle {
assert_ne!(offset, CURRENT_POSITION, "offset u64::MAX is reserved");
self.submit(file, offset, len, 1)
}
pub fn read_current(&self, file: Arc<File>, len: usize) -> ReadHandle {
self.submit(file, CURRENT_POSITION, len, 1)
}
pub fn read_at_direct(&self, file: Arc<File>, offset: u64, len: usize, align: usize) -> ReadHandle {
assert_ne!(offset, CURRENT_POSITION, "offset u64::MAX is reserved");
self.submit(file, offset, len, align)
}
fn submit(&self, file: Arc<File>, offset: u64, len: usize, align: usize) -> ReadHandle {
let id = self.next_id.fetch_add(1, Ordering::Relaxed);
assert_eq!(id & CANCEL_BIT, 0, "op id overflowed into the cancel bit");
let (done, rx) = oneshot::channel();
let shard = self.shard();
if offset != CURRENT_POSITION && offset > i64::MAX as u64 {
let _ = done.send(Err(io::Error::new(
io::ErrorKind::InvalidInput,
"offset exceeds i64::MAX (kernel loff_t is signed)",
)));
return ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: false,
state: HandleState::Inert,
};
}
if len > MAX_READ_LEN {
let _ = done.send(Err(io::Error::new(
io::ErrorKind::InvalidInput,
"read length exceeds MAX_RW_COUNT (2 GiB - 4 KiB); caller must chunk",
)));
return ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: false,
state: HandleState::Inert,
};
}
match aligned_geometry(offset, len, align) {
Some((kernel_offset, _, region_len))
if region_len <= MAX_READ_LEN
&& (offset == CURRENT_POSITION
|| kernel_offset
.checked_add(region_len as u64)
.is_some_and(|end| end <= i64::MAX as u64)) => {}
_ => {
let _ = done.send(Err(io::Error::new(
io::ErrorKind::InvalidInput,
"alignment must be a power of two, and the block-aligned range must fit MAX_RW_COUNT and end within i64::MAX",
)));
return ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: false,
state: HandleState::Inert,
};
}
}
match Arc::clone(&shard.sem).try_acquire_owned() {
Ok(permit) => {
if let Err(mpsc::SendError(msg)) = shard.tx.send(Msg::Read {
id,
file,
offset,
len,
done,
permit,
align,
}) {
if let Msg::Read { done, .. } = msg {
let _ = done.send(Err(io::Error::other("uring driver shut down")));
}
return ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: false,
state: HandleState::Inert,
};
}
shard.wake_efd.signal();
ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: true,
state: HandleState::Submitted {
wake: Arc::clone(&shard.wake_efd),
},
}
}
Err(TryAcquireError::NoPermits) => ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: true,
state: HandleState::WaitingPermit {
acquire: Box::pin(Arc::clone(&shard.sem).acquire_owned()),
file,
offset,
len,
align,
done,
wake: Arc::clone(&shard.wake_efd),
},
},
Err(TryAcquireError::Closed) => {
let _ = done.send(Err(io::Error::other("uring driver shut down")));
ReadHandle {
id,
rx,
tx: shard.tx.clone(),
finished: false,
cancel_on_drop: false,
state: HandleState::Inert,
}
}
}
}
pub fn stats(&self) -> StatsSnapshot {
let mut snap = StatsSnapshot::default();
for shard in &self.shards {
let s = &shard.stats;
snap.submitted += s.submitted.load(Ordering::SeqCst);
snap.delivered += s.delivered.load(Ordering::SeqCst);
snap.orphan_reclaimed += s.orphan_reclaimed.load(Ordering::SeqCst);
snap.in_flight += s.in_flight.load(Ordering::SeqCst);
snap.cancel_succeeded += s.cancel_succeeded.load(Ordering::SeqCst);
snap.cancel_not_found += s.cancel_not_found.load(Ordering::SeqCst);
snap.cancel_already += s.cancel_already.load(Ordering::SeqCst);
snap.cq_overflow += s.cq_overflow.load(Ordering::SeqCst);
snap.submit_errors += s.submit_errors.load(Ordering::SeqCst);
}
snap
}
#[cfg(feature = "fault-injection")]
pub fn test_inject_panic(&self) {
let shard = self.shard();
let _ = shard.tx.send(Msg::TestPanic);
shard.wake_efd.signal();
}
pub fn shutdown(mut self) -> StatsSnapshot {
for shard in &self.shards {
let _ = shard.tx.send(Msg::Shutdown);
shard.wake_efd.signal();
}
for shard in &mut self.shards {
shard.join();
}
let snap = self.stats();
if snap.in_flight != 0 {
eprintln!(
"uring-spike shutdown: {} ops still in flight (bounded-drain bailout on a hung device)",
snap.in_flight
);
}
snap
}
}
impl Drop for UringDriver {
fn drop(&mut self) {
for shard in &self.shards {
let _ = shard.tx.send(Msg::Shutdown);
shard.wake_efd.signal();
}
for shard in &mut self.shards {
shard.join();
}
}
}
fn probe_real_read(ring: &mut IoUring) -> io::Result<()> {
let pattern: Vec<u8> = (0..512u32).map(|i| (i * 7 + 13) as u8).collect();
let file = open_probe_file(&pattern)?;
let mut buf = vec![0u8; pattern.len()];
let sqe = opcode::Read::new(types::Fd(file.as_raw_fd()), buf.as_mut_ptr(), buf.len() as u32)
.offset(0)
.build()
.user_data(0xB0BE);
if unsafe { ring.submission().push(&sqe) }.is_err() {
return Err(io::Error::other("probe: submission queue full"));
}
let res = match drain_probe_cqe(ring) {
Ok(res) => res,
Err(e) => {
std::mem::forget(buf);
std::mem::forget(file);
return Err(e);
}
};
if res < 0 {
Err(io::Error::from_raw_os_error(-res))
} else if res as usize != pattern.len() || buf != pattern {
Err(io::Error::other("probe: read completed but data mismatched"))
} else {
Ok(())
}
}
fn open_probe_file(pattern: &[u8]) -> io::Result<File> {
let dir = std::env::temp_dir();
let c_dir = std::ffi::CString::new(dir.as_os_str().as_bytes()).map_err(|_| io::Error::other("probe dir path has NUL"))?;
let fd = unsafe { libc::open(c_dir.as_ptr(), libc::O_TMPFILE | libc::O_RDWR | libc::O_CLOEXEC, 0o600) };
if fd >= 0 {
let mut file = unsafe { File::from_raw_fd(fd) };
file.write_all(pattern)?;
return Ok(file);
}
open_probe_file_exclusive(&dir, pattern)
}
fn open_probe_file_exclusive(dir: &std::path::Path, pattern: &[u8]) -> io::Result<File> {
static SEQ: AtomicU64 = AtomicU64::new(0);
let nonce = SEQ.fetch_add(1, Ordering::Relaxed);
let path = dir.join(format!("uring-spike-probe-{}-{}", std::process::id(), nonce));
let c_path = std::ffi::CString::new(path.as_os_str().as_bytes()).map_err(|_| io::Error::other("probe path has NUL"))?;
let fd = unsafe {
libc::open(
c_path.as_ptr(),
libc::O_CREAT | libc::O_EXCL | libc::O_NOFOLLOW | libc::O_RDWR | libc::O_CLOEXEC,
0o600,
)
};
if fd < 0 {
return Err(io::Error::last_os_error());
}
let mut file = unsafe { File::from_raw_fd(fd) };
file.write_all(pattern)?;
unsafe {
libc::unlink(c_path.as_ptr());
}
Ok(file)
}
fn drain_probe_cqe(ring: &mut IoUring) -> io::Result<i32> {
const PROBE_TIMEOUT: Duration = Duration::from_secs(2);
let deadline = Instant::now() + PROBE_TIMEOUT;
let ext_arg = ring.params().is_feature_ext_arg();
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(io::Error::other("probe: no CQE within the bounded wait"));
}
let waited = if ext_arg {
let ts = types::Timespec::new().sec(remaining.as_secs()).nsec(remaining.subsec_nanos());
let args = types::SubmitArgs::new().timespec(&ts);
ring.submitter().submit_with_args(1, &args)
} else {
ring.submit_and_wait(1)
};
match waited {
Ok(_) => {}
Err(e) if e.raw_os_error() == Some(libc::EINTR) => {}
Err(e) if e.raw_os_error() == Some(libc::ETIME) => {}
Err(e) => return Err(e),
}
if let Some(cqe) = ring.completion().next() {
#[cfg(feature = "fault-injection")]
if std::env::var_os("RUSTFS_URING_FAULT_PROBE_DRAIN").is_some() {
return Err(io::Error::other("fault-injection: forced probe drain failure"));
}
return Ok(cqe.result());
}
}
}
struct DriverState {
ring: IoUring,
pending: HashMap<u64, Pending>,
backlog: VecDeque<io_uring::squeue::Entry>,
}
impl Drop for DriverState {
fn drop(&mut self) {
if std::thread::panicking() {
eprintln!(
"uring-spike driver thread panicked with {} ops in flight; \
aborting to avoid UAF of in-flight buffers",
self.pending.len()
);
std::process::abort();
}
}
}
fn file_len(file: &File) -> Option<u64> {
file.metadata().ok().map(|m| m.len())
}
fn deliver(p: &mut Pending) -> Vec<u8> {
let avail = p.nread.saturating_sub(p.head).min(p.want);
let start = p.pad + p.head;
if start != 0 && avail != 0 {
p.buf.copy_within(start..start + avail, 0);
}
p.buf.truncate(avail);
std::mem::take(&mut p.buf)
}
enum ReapStep {
Finish(io::Result<Vec<u8>>),
Resubmit(io_uring::squeue::Entry),
}
fn queue_cancel(backlog: &mut VecDeque<io_uring::squeue::Entry>, queued_cancels: &mut HashSet<u64>, id: u64) {
if queued_cancels.insert(id) {
backlog.push_back(opcode::AsyncCancel::new(id).build().user_data(id | CANCEL_BIT));
}
}
fn flush_backlog(ring: &mut IoUring, backlog: &mut VecDeque<io_uring::squeue::Entry>) {
let mut sq = ring.submission();
while let Some(sqe) = backlog.pop_front() {
if unsafe { sq.push(&sqe) }.is_err() {
backlog.push_front(sqe);
break;
}
}
}
fn submit_ring(
state: &mut DriverState,
stats: &DriverStats,
consecutive_submit_errors: &mut u32,
submit_error_logged: &mut bool,
shutting_down: &mut bool,
queued_cancels: &mut HashSet<u64>,
) {
flush_backlog(&mut state.ring, &mut state.backlog);
if state.ring.submission().is_empty() {
return;
}
match state.ring.submit() {
Ok(_) => *consecutive_submit_errors = 0,
Err(e) if matches!(e.raw_os_error(), Some(libc::EBUSY) | Some(libc::EINTR)) => *consecutive_submit_errors = 0,
Err(e) => {
stats.submit_errors.fetch_add(1, Ordering::SeqCst);
*consecutive_submit_errors += 1;
if !*submit_error_logged {
*submit_error_logged = true;
eprintln!("uring-spike driver: ring.submit() failed ({e}); retrying, will shut down if persistent");
}
if !*shutting_down && *consecutive_submit_errors >= MAX_CONSECUTIVE_SUBMIT_ERRORS {
eprintln!(
"uring-spike driver: {} consecutive submit failures; shutting down so callers fall back to the std backend",
*consecutive_submit_errors
);
*shutting_down = true;
let ids: Vec<u64> = state.pending.keys().copied().collect();
for id in ids {
queue_cancel(&mut state.backlog, queued_cancels, id);
}
}
}
}
}
fn drive(
ring: IoUring,
rx: mpsc::Receiver<Msg>,
stats: Arc<DriverStats>,
sem: Arc<Semaphore>,
cq_efd: EventFd,
wake_efd: Arc<EventFd>,
) {
let mut state = DriverState {
ring,
pending: HashMap::new(),
backlog: VecDeque::new(),
};
let mut shutting_down = false;
let mut drain_deadline: Option<Instant> = None;
let mut consecutive_submit_errors: u32 = 0;
let mut submit_error_logged = false;
let mut queued_cancels: HashSet<u64> = HashSet::new();
#[cfg(not(feature = "fault-injection"))]
let drain_timeout = DRAIN_TIMEOUT;
#[cfg(feature = "fault-injection")]
let drain_timeout = std::env::var("RUSTFS_URING_FAULT_DRAIN_TIMEOUT_MS")
.ok()
.and_then(|ms| ms.parse().ok())
.map(Duration::from_millis)
.unwrap_or(DRAIN_TIMEOUT);
#[cfg(feature = "fault-injection")]
let fault_stuck_drain = std::env::var_os("RUSTFS_URING_FAULT_STUCK_DRAIN").is_some();
loop {
let heartbeat = if shutting_down || !state.pending.is_empty() {
LOOP_HEARTBEAT
} else {
IDLE_HEARTBEAT
};
wait_for_events(&cq_efd, &wake_efd, heartbeat);
cq_efd.drain();
wake_efd.drain();
loop {
let msg = match rx.try_recv() {
Ok(m) => m,
Err(TryRecvError::Empty) => break,
Err(TryRecvError::Disconnected) => {
shutting_down = true;
break;
}
};
match msg {
Msg::Read {
id,
file,
offset,
len,
done,
permit,
align,
} => {
if shutting_down {
let _ = done.send(Err(io::Error::other("uring driver shutting down")));
drop(permit);
continue;
}
let (kernel_offset, head, region_len) =
aligned_geometry(offset, len, align).expect("submit validated the geometry");
let cap = match region_len.checked_add(align - 1) {
Some(cap) => cap,
None => {
let _ = done.send(Err(io::Error::other("aligned O_DIRECT allocation size overflow")));
drop(permit);
continue;
}
};
let buf = vec![0u8; cap];
let pad = buf.as_ptr().align_offset(align);
if pad == usize::MAX || pad.checked_add(region_len).is_none_or(|end| end > buf.len()) {
let _ = done.send(Err(io::Error::other("could not align O_DIRECT read buffer")));
drop(permit);
continue;
}
state.pending.insert(
id,
Pending {
buf,
file,
done: Some(done),
offset: kernel_offset,
nread: 0,
_permit: permit,
pad,
head,
want: len,
region_len,
align,
transient_retries: 0,
},
);
let sqe = state.pending.get(&id).expect("just inserted").read_sqe(id);
stats.submitted.fetch_add(1, Ordering::SeqCst);
stats.in_flight.fetch_add(1, Ordering::SeqCst);
state.backlog.push_back(sqe);
}
Msg::Cancel { id } => {
if state.pending.contains_key(&id) {
queue_cancel(&mut state.backlog, &mut queued_cancels, id);
}
}
Msg::Shutdown => {
shutting_down = true;
let ids: Vec<u64> = state.pending.keys().copied().collect();
for id in ids {
queue_cancel(&mut state.backlog, &mut queued_cancels, id);
}
}
#[cfg(feature = "fault-injection")]
Msg::TestPanic => {
panic!(
"fault-injection: driver thread panic requested with {} ops in flight",
state.pending.len()
);
}
}
}
submit_ring(
&mut state,
&stats,
&mut consecutive_submit_errors,
&mut submit_error_logged,
&mut shutting_down,
&mut queued_cancels,
);
while let Some(cqe) = state.ring.completion().next() {
let ud = cqe.user_data();
if ud & CANCEL_BIT != 0 {
match cqe.result() {
0 => stats.cancel_succeeded.fetch_add(1, Ordering::SeqCst),
r if r == -libc::ENOENT => stats.cancel_not_found.fetch_add(1, Ordering::SeqCst),
r if r == -libc::EALREADY => stats.cancel_already.fetch_add(1, Ordering::SeqCst),
_ => 0,
};
continue;
}
#[cfg(feature = "fault-injection")]
if fault_stuck_drain && state.pending.contains_key(&ud) {
continue;
}
let res = cqe.result();
if !state.pending.contains_key(&ud) {
continue;
}
let step = {
let p = state.pending.get_mut(&ud).expect("checked above");
if res < 0 {
let err = -res;
let transient = err == libc::EINTR || err == libc::EAGAIN;
if transient
&& p.offset != CURRENT_POSITION
&& p.nread < p.region_len
&& p.transient_retries < MAX_TRANSIENT_RETRIES
{
p.transient_retries += 1;
ReapStep::Resubmit(p.read_sqe(ud))
} else {
ReapStep::Finish(Err(io::Error::from_raw_os_error(err)))
}
} else if res == 0 {
ReapStep::Finish(Ok(deliver(p)))
} else {
p.nread += res as usize;
p.transient_retries = 0;
let is_stream = p.offset == CURRENT_POSITION;
let covered = p.nread >= p.head + p.want;
if is_stream || covered || p.nread >= p.region_len {
ReapStep::Finish(Ok(deliver(p)))
} else if p.align > 1 && !p.nread.is_multiple_of(p.align) {
match file_len(&p.file) {
Some(len) if p.offset + p.nread as u64 >= len => ReapStep::Finish(Ok(deliver(p))),
Some(_) => ReapStep::Finish(Err(io::Error::other(
"io_uring O_DIRECT: non-block-aligned short read before EOF",
))),
None => ReapStep::Finish(Ok(deliver(p))),
}
} else {
ReapStep::Resubmit(p.read_sqe(ud))
}
}
};
match step {
ReapStep::Finish(outcome) => {
let mut p = state.pending.remove(&ud).expect("checked above");
match p.done.take().expect("done sender set at submit").send(outcome) {
Ok(()) => stats.delivered.fetch_add(1, Ordering::SeqCst),
Err(_) => stats.orphan_reclaimed.fetch_add(1, Ordering::SeqCst),
};
stats.in_flight.fetch_sub(1, Ordering::SeqCst);
queued_cancels.remove(&ud);
}
ReapStep::Resubmit(sqe) => state.backlog.push_back(sqe),
}
}
if !state.backlog.is_empty() {
submit_ring(
&mut state,
&stats,
&mut consecutive_submit_errors,
&mut submit_error_logged,
&mut shutting_down,
&mut queued_cancels,
);
}
let overflow = state.ring.completion().overflow();
if overflow != 0 {
stats.cq_overflow.store(overflow as u64, Ordering::SeqCst);
eprintln!("uring-spike driver: CQ overflow = {overflow}; CQEs buffered (NODROP), not lost — backpressure warning");
}
if shutting_down {
if state.pending.is_empty() && state.backlog.is_empty() {
sem.close();
return; }
let deadline = *drain_deadline.get_or_insert_with(|| Instant::now() + drain_timeout);
if Instant::now() >= deadline {
eprintln!(
"uring-spike driver: bounded drain timed out with {} ops still in flight; \
leaking ring + buffers to stay memory-safe",
state.pending.len()
);
for p in state.pending.values_mut() {
if let Some(tx) = p.done.take() {
let _ = tx.send(Err(io::Error::other("uring driver leaked op on bounded-drain timeout")));
}
}
sem.close();
std::mem::forget(cq_efd);
std::mem::forget(state);
return;
}
}
}
}