use super::*;
use crate::{BLOB_PAGE_SIZE, BufMut as _, IoBufMut};
use commonware_utils::NZUsize;
use futures::stream::{self, BoxStream};
const NR_EVENTS: usize = 256;
const IOCBS_PER_SUBMIT: usize = 32;
struct Read {
index: usize,
offset: u64,
len: usize,
}
type Completion = Result<(usize, IoBufsMut), Error>;
pub(super) fn direct(file: &Shared) -> Option<&File> {
crate::storage::open_direct(&file.direct, file)
}
fn superset(offset: u64, len: usize) -> Option<(u64, usize, usize)> {
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let within = offset % u64::from(BLOB_PAGE_SIZE);
let skip = usize::try_from(within).expect("an offset within a block must fit in usize");
let aligned_len = skip.checked_add(len)?.checked_next_multiple_of(block)?;
Some((offset - within, skip, aligned_len))
}
pub(super) fn read_direct(file: &Shared, offset: u64, len: usize) -> Option<IoBufMut> {
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let (aligned_offset, skip, aligned_len) = superset(offset, len)?;
let end = skip + len;
aligned_offset
.checked_add(Widen::widen(aligned_len))
.filter(|superset_end| *superset_end <= i64::MAX as u64)?;
let direct = direct(file)?;
let mut buf = IoBufMut::try_with_alignment(aligned_len, NZUsize!(block))?;
let mut filled = 0;
while filled < end {
let position = libc::off_t::try_from(aligned_offset + Widen::widen(filled)).ok()?;
let n = unsafe {
libc::pread(
direct.as_raw_fd(),
buf.as_mut_ptr().add(filled).cast(),
aligned_len - filled,
position,
)
};
if n < 0 {
if std::io::Error::last_os_error().kind() == std::io::ErrorKind::Interrupted {
continue;
}
return None;
}
let n = n as usize;
filled += n;
if filled < end && (n == 0 || !n.is_multiple_of(block)) {
return None;
}
}
unsafe { buf.set_len(end) };
buf.advance(skip);
Some(buf)
}
#[repr(C)]
#[derive(Default, Clone, Copy)]
struct Iocb {
aio_data: u64,
aio_key: u32,
aio_rw_flags: u32,
aio_lio_opcode: u16,
aio_reqprio: i16,
aio_fildes: u32,
aio_buf: u64,
aio_nbytes: u64,
aio_offset: i64,
aio_reserved2: u64,
aio_flags: u32,
aio_resfd: u32,
}
#[repr(C)]
#[derive(Default, Clone, Copy)]
struct IoEvent {
data: u64,
obj: u64,
res: i64,
res2: i64,
}
const IOCB_CMD_PREAD: u16 = 0;
struct Context(libc::c_ulong);
impl Drop for Context {
fn drop(&mut self) {
let result = unsafe { libc::syscall(libc::SYS_io_destroy, self.0) };
if result != 0 {
std::process::abort();
}
}
}
static CONTEXTS: Mutex<Vec<Context>> = Mutex::new(Vec::new());
fn take_context() -> Option<Context> {
if let Some(ctx) = CONTEXTS.lock().pop() {
return Some(ctx);
}
let mut ctx: libc::c_ulong = 0;
let r = unsafe { libc::syscall(libc::SYS_io_setup, NR_EVENTS as libc::c_ulong, &mut ctx) };
(r == 0).then(|| Context(ctx))
}
fn read_positioned(file: &Shared, read: &Read) -> Completion {
let mut buf = IoBufMut::with_capacity(read.len);
unsafe { buf.set_len(read.len) };
Blob::read_exact_at(Cache::Disabled, file, buf.as_mut(), read.offset)?;
Ok((read.index, buf.into()))
}
fn submit(
file: &Shared,
direct: &File,
ctx: Context,
reads: Vec<Read>,
tx: &tokio::sync::mpsc::UnboundedSender<Completion>,
) -> Result<Context, Error> {
let fd = u32::try_from(direct.as_raw_fd()).expect("an open descriptor is non-negative");
let n = reads.len();
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let mut spans = Vec::with_capacity(n);
let mut slab_len = 0usize;
for read in &reads {
let (aligned_offset, skip, aligned_len) =
superset(read.offset, read.len).ok_or(Error::OffsetOverflow)?;
spans.push((aligned_offset, skip, slab_len, aligned_len));
slab_len = slab_len
.checked_add(aligned_len)
.ok_or(Error::OffsetOverflow)?;
}
let mut slab =
IoBufMut::try_with_alignment(slab_len, NZUsize!(block)).ok_or(Error::OffsetOverflow)?;
let slab_ptr = slab.chunk_mut().as_mut_ptr();
let active = ctx;
let mut iocbs: Vec<Iocb> = spans
.iter()
.enumerate()
.map(|(i, &(aligned_offset, _, start, aligned_len))| Iocb {
aio_data: Widen::widen(i),
aio_lio_opcode: IOCB_CMD_PREAD,
aio_fildes: fd,
aio_buf: Widen::widen(unsafe { slab_ptr.add(start) }.expose_provenance()),
aio_nbytes: Widen::widen(aligned_len),
aio_offset: aligned_offset as i64,
..Iocb::default()
})
.collect();
let ptrs: Vec<*mut Iocb> = iocbs.iter_mut().map(|iocb| iocb as *mut Iocb).collect();
let deliver = |event: &IoEvent| -> Result<(), Error> {
let i = event.data as usize;
let read = &reads[i];
let (_, skip, start, _) = spans[i];
let item = if event.res < 0 || (event.res as usize) < skip + read.len {
read_positioned(file, read)?
} else {
let bytes = unsafe { std::slice::from_raw_parts(slab_ptr.add(start + skip), read.len) };
let mut buf = IoBufMut::with_capacity(read.len);
buf.put_slice(bytes);
#[cfg(test)]
file.test.direct_reads.fetch_add(1, Ordering::Relaxed);
(read.index, buf.into())
};
let _ = tx.send(Ok(item));
Ok(())
};
let mut events = vec![IoEvent::default(); n];
let mut next = 0;
let mut accepted = 0;
let mut completed = 0;
while next < n {
let count = (n - next).min(IOCBS_PER_SUBMIT);
let r = unsafe {
libc::syscall(
libc::SYS_io_submit,
active.0,
count as libc::c_long,
ptrs.as_ptr().add(next),
)
};
if r < 0 {
let _ = tx.send(Ok(read_positioned(file, &reads[next])?));
next += 1;
continue;
}
next += r as usize;
accepted += r as usize;
if next < n {
let got = reap(&active, &mut events[..accepted - completed], false)?;
events[..got].iter().try_for_each(&deliver)?;
completed += got;
}
}
while completed < accepted {
let got = reap(&active, &mut events[..accepted - completed], true)?;
events[..got].iter().try_for_each(&deliver)?;
completed += got;
}
Ok(active)
}
fn reap(ctx: &Context, events: &mut [IoEvent], wait: bool) -> Result<usize, Error> {
if events.is_empty() {
return Ok(0);
}
let zero = libc::timespec {
tv_sec: 0,
tv_nsec: 0,
};
let timeout: *const libc::timespec = if wait { std::ptr::null() } else { &zero };
loop {
let got = unsafe {
libc::syscall(
libc::SYS_io_getevents,
ctx.0,
wait as libc::c_long,
events.len() as libc::c_long,
events.as_mut_ptr(),
timeout,
)
};
if got >= 0 {
return Ok(got as usize);
}
let err = std::io::Error::last_os_error();
if err.kind() != std::io::ErrorKind::Interrupted {
return Err(err.into());
}
}
}
fn read_positioned_each(
file: &Arc<Shared>,
reads: Vec<Read>,
tx: &tokio::sync::mpsc::UnboundedSender<Completion>,
) {
for read in reads {
let (file, tx) = (file.clone(), tx.clone());
task::spawn_blocking(move || {
let _ = tx.send(read_positioned(&file, &read));
});
}
}
fn run(file: &Arc<Shared>, reads: Vec<Read>, tx: &tokio::sync::mpsc::UnboundedSender<Completion>) {
let Some(direct) = direct(file) else {
return read_positioned_each(file, reads, tx);
};
let Some(ctx) = take_context() else {
return read_positioned_each(file, reads, tx);
};
match submit(file, direct, ctx, reads, tx) {
Ok(ctx) => CONTEXTS.lock().push(ctx),
Err(err) => {
let _ = tx.send(Err(err));
}
}
}
pub(super) fn read_many(
file: &Arc<Shared>,
data_offset: u64,
ranges: &[(u64, usize)],
) -> BoxStream<'static, Completion> {
let mut reads = Vec::with_capacity(ranges.len());
for (index, &(offset, len)) in ranges.iter().enumerate() {
let Some(offset) = physical_offset(data_offset, offset, len) else {
return stream::iter(vec![Err(Error::OffsetOverflow)]).boxed();
};
reads.push(Read { index, offset, len });
}
let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
let mut reads = reads.into_iter().peekable();
while reads.peek().is_some() {
let chunk: Vec<_> = reads.by_ref().take(NR_EVENTS).collect();
let (file, tx) = (file.clone(), tx.clone());
task::spawn_blocking(move || run(&file, chunk, &tx));
}
stream::unfold((rx, ranges.len()), |(mut rx, remaining)| async move {
if remaining == 0 {
return None;
}
let item = rx.recv().await.unwrap_or(Err(Error::ReadFailed));
let remaining = if item.is_ok() { remaining - 1 } else { 0 };
Some((item, (rx, remaining)))
})
.boxed()
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
Blob as _, Storage as _,
storage::{
Layout,
tokio::{Config, Storage},
},
};
use futures::TryStreamExt as _;
use std::{
path::PathBuf,
sync::atomic::{AtomicBool, Ordering},
time::Duration,
};
async fn direct_blob(label: &str, data: Vec<u8>) -> (Storage, Blob, PathBuf) {
let directory =
std::env::temp_dir().join(format!("storage_tokio_aio_{label}_{}", std::process::id()));
let storage = Storage::new(Config::new(directory.clone(), Layout::ALL));
let (blob, _) = storage.open("partition", b"blob").await.unwrap();
blob.write_at(0, data, WriteOptions::SYNC).await.unwrap();
assert!(
direct(&blob.shared).is_some(),
"temporary directory must support O_DIRECT"
);
(storage, blob, directory)
}
async fn remove(storage: Storage, blob: Blob, directory: PathBuf) {
drop(blob);
storage.remove("partition", None).await.unwrap();
drop(storage);
std::fs::remove_dir_all(directory).unwrap();
}
#[tokio::test]
async fn test_failed_submission_keeps_slab_until_reads_complete() {
const LEN: usize = 64 << 10;
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let (storage, blob, directory) = direct_blob("failed_submission", vec![0xAB; LEN]).await;
let (tx, _rx) = tokio::sync::mpsc::unbounded_channel();
for _ in 0..20 {
let reads = vec![
Read {
index: 0,
offset: 0,
len: LEN,
},
Read {
index: 1,
offset: 1 << 63,
len: block,
},
];
let ctx = take_context().unwrap();
let direct = direct(&blob.shared).unwrap();
assert!(submit(&blob.shared, direct, ctx, reads, &tx).is_err());
let canary = IoBufMut::zeroed_with_alignment(LEN + block, NZUsize!(block));
std::thread::sleep(Duration::from_millis(10));
assert!(!canary.as_ref().contains(&0xAB));
}
remove(storage, blob, directory).await;
}
#[tokio::test]
async fn test_submission_beyond_one_slab_fails() {
let (storage, blob, directory) = direct_blob("slab_overflow", vec![0; 4096]).await;
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let len = 1usize << (usize::BITS - 2);
let reads = (0..2)
.map(|index| Read {
index,
offset: 0,
len,
})
.collect();
let ctx = take_context().unwrap();
let direct = direct(&blob.shared).unwrap();
let result = submit(&blob.shared, direct, ctx, reads, &tx);
assert!(matches!(result, Err(Error::OffsetOverflow)));
drop(tx);
assert!(rx.recv().await.is_none());
remove(storage, blob, directory).await;
}
#[tokio::test]
async fn test_submission_survives_signals() {
extern "C" fn ignore(_: libc::c_int) {}
let previous = unsafe {
let mut action: libc::sigaction = std::mem::zeroed();
action.sa_sigaction = ignore as extern "C" fn(libc::c_int) as usize;
let mut previous: libc::sigaction = std::mem::zeroed();
assert_eq!(libc::sigaction(libc::SIGUSR1, &action, &mut previous), 0);
previous
};
const LEN: usize = 64 << 20;
let (storage, blob, directory) = direct_blob("signals", vec![0xAB; LEN]).await;
let reads = vec![Read {
index: 0,
offset: blob.data_offset,
len: LEN,
}];
let shared = blob.shared.clone();
let done = Arc::new(AtomicBool::new(false));
let (thread_tx, thread_rx) = std::sync::mpsc::channel();
let submitter = {
let done = done.clone();
std::thread::spawn(move || {
thread_tx.send(unsafe { libc::pthread_self() }).unwrap();
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
let ctx = take_context().unwrap();
let direct = direct(&shared).unwrap();
let result = submit(&shared, direct, ctx, reads, &tx);
done.store(true, Ordering::Release);
(result.is_ok(), rx.try_recv().ok(), rx.try_recv().is_err())
})
};
let thread = thread_rx.recv().unwrap();
while !done.load(Ordering::Acquire) {
unsafe { libc::pthread_kill(thread, libc::SIGUSR1) };
std::thread::sleep(Duration::from_micros(10));
}
let (submitted, first, drained) = submitter.join().unwrap();
unsafe { libc::sigaction(libc::SIGUSR1, &previous, std::ptr::null_mut()) };
assert!(submitted);
assert!(drained);
let Some(Ok((0, bufs))) = first else {
panic!("expected the read");
};
assert_eq!(bufs.coalesce().as_ref(), vec![0xAB; LEN].as_slice());
remove(storage, blob, directory).await;
}
#[tokio::test]
async fn test_direct_open_retries_after_descriptor_limit() {
let directory = std::env::temp_dir().join(format!(
"storage_tokio_aio_retry_open_{}",
std::process::id()
));
let storage = Storage::new(Config::new(directory.clone(), Layout::ALL));
let (blob, _) = storage.open("partition", b"blob").await.unwrap();
blob.write_at(0, vec![0; Widen::widen(BLOB_PAGE_SIZE)], WriteOptions::SYNC)
.await
.unwrap();
let denied = unsafe {
let mut previous: libc::rlimit = std::mem::zeroed();
assert_eq!(libc::getrlimit(libc::RLIMIT_NOFILE, &mut previous), 0);
let none = libc::rlimit {
rlim_cur: 0,
rlim_max: previous.rlim_max,
};
assert_eq!(libc::setrlimit(libc::RLIMIT_NOFILE, &none), 0);
let denied = direct(&blob.shared);
assert_eq!(libc::setrlimit(libc::RLIMIT_NOFILE, &previous), 0);
denied.is_none()
};
assert!(denied);
assert!(blob.shared.direct.get().is_none());
let opened = direct(&blob.shared).is_some();
assert_eq!(opened, matches!(blob.shared.direct.get(), Some(Some(_))));
assert!(blob.shared.direct.get().is_some());
remove(storage, blob, directory).await;
}
#[tokio::test]
async fn test_read_many_serves_more_ranges_than_one_submission() {
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let data: Vec<u8> = (0..8 * block).map(|i| (i % 251) as u8).collect();
let (storage, blob, directory) = direct_blob("many_submissions", data.clone()).await;
let ranges: Vec<(u64, usize)> = (0..2 * NR_EVENTS + 1)
.map(|i| (Widen::widen((i % 16) * 2048), 2048))
.collect();
let before = blob.shared.test.direct_reads.load(Ordering::Relaxed);
let mut served: Vec<(usize, IoBufsMut)> =
read_many(&blob.shared, blob.data_offset, &ranges)
.try_collect()
.await
.unwrap();
assert_eq!(
blob.shared.test.direct_reads.load(Ordering::Relaxed) - before,
ranges.len()
);
served.sort_by_key(|(index, _)| *index);
assert!(served.iter().map(|(index, _)| *index).eq(0..ranges.len()));
for (index, bufs) in served {
let (offset, len) = ranges[index];
let offset = usize::try_from(offset).unwrap();
assert_eq!(bufs.coalesce().as_ref(), &data[offset..offset + len]);
}
remove(storage, blob, directory).await;
}
#[tokio::test]
async fn test_reads_without_context_are_served_individually() {
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let data: Vec<u8> = (0..3 * block).map(|i| (i % 251) as u8).collect();
let (storage, blob, directory) = direct_blob("without_context", data.clone()).await;
let ranges = [
(10u64, 5000usize),
(u64::from(BLOB_PAGE_SIZE) - 1, 2),
(data.len() as u64, 1),
];
let reads = ranges
.iter()
.enumerate()
.map(|(index, &(offset, len))| Read {
index,
offset: offset + blob.data_offset,
len,
})
.collect();
let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
read_positioned_each(&blob.shared, reads, &tx);
drop(tx);
let mut served = Vec::new();
let mut failed = 0;
while let Some(item) = rx.recv().await {
match item {
Ok(item) => served.push(item),
Err(_) => failed += 1,
}
}
assert_eq!(failed, 1);
served.sort_by_key(|(index, _)| *index);
assert_eq!(served.len(), 2);
for ((offset, len), (index, bufs)) in ranges.iter().zip(served) {
let offset = *offset as usize;
assert_eq!(
bufs.coalesce().as_ref(),
&data[offset..offset + len],
"{index}"
);
}
remove(storage, blob, directory).await;
}
#[commonware_macros::test_group("miri")]
#[test]
fn test_read_direct_covers_the_aligned_superset() {
let block: usize = Widen::widen(BLOB_PAGE_SIZE);
let directory = std::env::temp_dir().join(format!(
"storage_tokio_aio_read_direct_{}",
std::process::id()
));
let storage = Storage::new(Config::new(directory.clone(), Layout::ALL));
let data: Vec<u8> = (0..3 * block).map(|i| (i % 251) as u8).collect();
let open = |name: &[u8], len: usize| {
let path = directory.join(hex(name));
std::fs::write(&path, &data[..len]).unwrap();
let (generation, _, _) = storage.pending.admit("partition", name, false).unwrap();
let blob = Blob::new(
File::open(&path).unwrap(),
0,
storage.hold.clone(),
generation,
);
assert!(
blob.shared
.direct
.set(Some(File::open(&path).unwrap()))
.is_ok()
);
blob
};
let served = |blob: &Blob, offset: u64, len: usize| {
let attempts = if cfg!(miri) { 64 } else { 1 };
(0..attempts)
.find_map(|_| read_direct(&blob.shared, offset, len))
.expect("the file must serve the range")
};
let unaligned = open(b"unaligned", 3 * block - 50);
for (offset, len) in [(0, block), (100, 5000), (2 * block + 10, block - 60)] {
let buf = served(&unaligned, offset as u64, len);
let skip = offset % block;
assert_eq!(buf.as_ref(), &data[offset..offset + len]);
assert_eq!(buf.as_ref().as_ptr().addr() % block, skip);
assert_eq!(buf.capacity(), (skip + len).next_multiple_of(block) - skip);
}
for (offset, len) in [
(2 * block as u64 + 10, block),
(3 * block as u64, 1),
(0, usize::MAX - 10),
(i64::MAX as u64 - 10, 1),
] {
assert!(read_direct(&unaligned.shared, offset, len).is_none());
}
let aligned = open(b"aligned", 2 * block);
assert!(read_direct(&aligned.shared, 100, 2 * block).is_none());
let buf = served(&aligned, 100, block);
assert_eq!(buf.as_ref(), &data[100..100 + block]);
drop(unaligned);
drop(aligned);
drop(storage);
std::fs::remove_dir_all(directory).unwrap();
}
}