use std::fs::File;
use std::io;
pub fn platform_read_at(f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
platform_read_at_impl(f, buf, off)
}
#[cfg(unix)]
fn platform_read_at_impl(f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
use std::os::unix::fs::FileExt;
f.read_at(buf, off)
}
#[cfg(windows)]
fn platform_read_at_impl(f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
use std::os::windows::fs::FileExt;
f.seek_read(buf, off)
}
#[cfg(not(any(unix, windows)))]
fn platform_read_at_impl(f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
let mut f2 = f.try_clone()?;
use std::io::{Read, Seek, SeekFrom};
f2.seek(SeekFrom::Start(off))?;
f2.read(buf)
}
pub fn platform_read_exact_at(f: &File, buf: &mut [u8], off: u64) -> io::Result<()> {
let mut done = 0usize;
while done < buf.len() {
let n = platform_read_at(f, &mut buf[done..], off + done as u64)?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"read_exact_at: eof inside buffer",
));
}
done += n;
}
Ok(())
}
pub struct PreadvBackend;
impl IoBackend for PreadvBackend {
fn read_exact_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<()> {
platform_read_exact_at(f, buf, off)
}
fn read_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
platform_read_at(f, buf, off)
}
}
pub trait IoBackend: Send + Sync + 'static {
fn read_exact_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<()>;
fn read_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<usize>;
fn read_many(&self, f: &File, reqs: &mut [(u64, &mut [u8])]) -> io::Result<()> {
for (off, buf) in reqs.iter_mut() {
self.read_exact_at(f, buf, *off)?;
}
Ok(())
}
}
pub fn default_backend() -> Box<dyn IoBackend> {
Box::new(PreadvBackend)
}
#[cfg(target_os = "linux")]
pub struct UringBackend;
#[cfg(target_os = "linux")]
pub struct UringBatchBackend;
#[cfg(target_os = "linux")]
impl IoBackend for UringBatchBackend {
fn read_exact_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<()> {
platform_read_exact_at(f, buf, off)
}
fn read_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
platform_read_at(f, buf, off)
}
fn read_many(&self, f: &File, reqs: &mut [(u64, &mut [u8])]) -> io::Result<()> {
let _ = uring_batch_read(f, reqs)?;
Ok(())
}
}
#[cfg(target_os = "linux")]
fn uring_batch_read(f: &File, reqs: &mut [(u64, &mut [u8])]) -> io::Result<usize> {
const DEPTH: u32 = 256;
use std::os::fd::AsRawFd;
URING.with(|sl| {
let mut r = sl.borrow_mut();
if r.is_none() {
let ring = io_uring::IoUring::new(DEPTH)
.map_err(|e| io::Error::other(format!("IoUring::new: {e}")))?;
*r = Some(ring);
}
let ring = r.as_mut().unwrap();
let fd = f.as_raw_fd();
let mut total = 0usize;
for chunk in reqs.chunks_mut(DEPTH as usize) {
let expect: Vec<(usize, usize)> = chunk
.iter()
.enumerate()
.map(|(i, (_, b))| (i, b.len()))
.collect();
for (i, (off, buf)) in chunk.iter_mut().enumerate() {
let sqe = io_uring::opcode::Read::new(
io_uring::types::Fd(fd),
buf.as_mut_ptr(),
buf.len() as u32,
)
.offset(*off);
let sqe = sqe.build().user_data(i as u64);
unsafe {
ring.submission()
.push(&sqe)
.map_err(|e| io::Error::other(format!("SQ push: {e}")))?;
}
}
let n = chunk.len();
ring.submit_and_wait(n)
.map_err(|e| io::Error::other(format!("submit_and_wait: {e}")))?;
for _ in 0..n {
let cqe = ring
.completion()
.next()
.ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "no cqe"))?;
let res = cqe.result();
let idx = cqe.user_data() as usize;
let elen = expect.get(idx).map(|x| x.1).unwrap_or(0);
if res < 0 {
return Err(io::Error::from_raw_os_error(-res));
}
if (res as usize) != elen {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
format!("read_many: short read {res} < {elen}"),
));
}
total += res as usize;
}
}
Ok(total)
})
}
#[cfg(target_os = "linux")]
thread_local! {
static URING: std::cell::RefCell<Option<io_uring::IoUring>> =
const { std::cell::RefCell::new(None) };
}
#[cfg(target_os = "linux")]
fn submit_uring_read(f: &std::fs::File, buf: &mut [u8], off: u64) -> io::Result<usize> {
use std::os::fd::AsRawFd;
URING.with(|sl| {
let mut r = sl.borrow_mut();
if r.is_none() {
let ring = io_uring::IoUring::new(256)
.map_err(|e| io::Error::other(format!("IoUring::new: {e}")))?;
*r = Some(ring);
}
let ring = r.as_mut().unwrap();
let fd = f.as_raw_fd();
let sqe = io_uring::opcode::Read::new(
io_uring::types::Fd(fd),
buf.as_mut_ptr(),
buf.len() as u32,
)
.offset(off);
let sqe = sqe.build().user_data(0);
unsafe {
ring.submission()
.push(&sqe)
.map_err(|e| io::Error::other(format!("SQ push: {e}")))?;
}
ring.submit_and_wait(1)
.map_err(|e| io::Error::other(format!("submit_and_wait: {e}")))?;
let cqe = ring
.completion()
.next()
.ok_or_else(|| io::Error::new(io::ErrorKind::UnexpectedEof, "no cqe"))?;
let res = cqe.result();
if res < 0 {
return Err(io::Error::from_raw_os_error(-res));
}
Ok(res as usize)
})
}
#[cfg(target_os = "linux")]
impl IoBackend for UringBackend {
fn read_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<usize> {
submit_uring_read(f, buf, off)
}
fn read_exact_at(&self, f: &File, buf: &mut [u8], off: u64) -> io::Result<()> {
let mut done = 0usize;
while done < buf.len() {
let n = self.read_at(f, &mut buf[done..], off + done as u64)?;
if n == 0 {
return Err(io::Error::new(
io::ErrorKind::UnexpectedEof,
"read_exact_at: eof inside buffer",
));
}
done += n;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn preadv_roundtrip() {
let dir = std::env::temp_dir().join("engramdb-backend-test");
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
let p = dir.join("t.bin");
std::fs::write(&p, (0u8..255).collect::<Vec<u8>>()).unwrap();
let f = std::fs::File::open(&p).unwrap();
let b = PreadvBackend;
let mut buf = [0u8; 4];
b.read_exact_at(&f, &mut buf, 10).unwrap();
assert_eq!(&buf, &[10, 11, 12, 13]);
let mut p2 = [0u8; 2];
let n = b.read_at(&f, &mut p2, 253).unwrap();
assert_eq!(n, 2);
assert_eq!(&p2, &[253, 254]);
let _ = std::fs::remove_dir_all(&dir);
}
#[test]
fn default_backend_constructs() {
let _ = default_backend();
}
#[cfg(target_os = "linux")]
#[test]
fn uring_roundtrip_and_semantics() {
let dir = std::env::temp_dir().join("engramdb-uring-test");
let _ = std::fs::remove_dir_all(&dir);
std::fs::create_dir_all(&dir).unwrap();
let p = dir.join("t.bin");
std::fs::write(&p, (0usize..512).map(|i| i as u8).collect::<Vec<u8>>()).unwrap();
let f = std::fs::File::open(&p).unwrap();
let b = UringBackend;
let mut buf = [0u8; 4];
b.read_exact_at(&f, &mut buf, 10).unwrap();
assert_eq!(&buf, &[10, 11, 12, 13]);
let mut p2 = [0u8; 2];
let n = b.read_at(&f, &mut p2, 253).unwrap();
assert_eq!(n, 2);
assert_eq!(&p2, &[253, 254]);
let mut p3 = [0u8; 4];
let n3 = b.read_at(&f, &mut p3, 512).unwrap();
assert_eq!(n3, 0);
let r = b.read_exact_at(&f, &mut p3, 511);
assert!(r.is_err());
let _ = std::fs::remove_dir_all(&dir);
}
}