#![allow(clippy::arithmetic_side_effects)]
use {
crate::{
FileInfo, FileSize, IoSize,
file_io::FileCreator,
io_uring::{
IO_PRIO_BE_HIGHEST,
memory::{IoBufferChunk, PageAlignedMemory},
sqpoll,
},
},
agave_io_uring::{Completion, FixedSlab, Ring, RingAccess, RingOp},
core::slice,
io_uring::{IoUring, opcode, squeue, types},
libc::{O_CREAT, O_DIRECT, O_NOATIME, O_NOFOLLOW, O_RDWR, O_TRUNC},
smallvec::SmallVec,
std::{
collections::VecDeque,
ffi::CString,
fs::File,
io::{self, Read},
mem,
os::fd::{AsRawFd, BorrowedFd, FromRawFd as _, IntoRawFd as _, RawFd},
path::PathBuf,
pin::Pin,
sync::Arc,
time::Duration,
},
};
pub const DEFAULT_WRITE_SIZE: IoSize = 512 * 1024;
const DIRECT_IO_WRITE_LEN_ALIGNMENT: IoSize = 4096;
const DEFAULT_STATUS_FLAGS: libc::c_int = O_NOATIME;
type BacklogVec = SmallVec<[PendingWrite; 8 * 1024 * 1024 / DEFAULT_WRITE_SIZE as usize]>;
const MAX_OPEN_FILES: usize = 512;
const DEFAULT_MAX_IOWQ_WORKERS: u32 = 4;
const CHECK_PROGRESS_AFTER_SUBMIT_TIMEOUT: Option<Duration> = Some(Duration::from_millis(10));
pub struct IoUringFileCreatorBuilder<'sp> {
write_capacity: IoSize,
max_iowq_workers: u32,
ring_squeue_size: Option<u32>,
shared_sqpoll_fd: Option<BorrowedFd<'sp>>,
register_buffer: bool,
write_with_direct_io: bool,
}
impl<'sp> IoUringFileCreatorBuilder<'sp> {
pub fn new() -> Self {
Self {
write_capacity: DEFAULT_WRITE_SIZE,
max_iowq_workers: DEFAULT_MAX_IOWQ_WORKERS,
ring_squeue_size: None,
shared_sqpoll_fd: None,
register_buffer: true,
write_with_direct_io: false,
}
}
#[cfg(test)]
pub fn write_capacity(mut self, write_capacity: IoSize) -> Self {
self.write_capacity = write_capacity;
self
}
pub fn use_registered_buffers(mut self, register_buffers: bool) -> Self {
self.register_buffer = register_buffers;
self
}
pub fn write_with_direct_io(mut self, enable_direct_io: bool) -> Self {
self.write_with_direct_io = enable_direct_io;
self
}
pub fn shared_sqpoll(mut self, shared_sqpoll_fd: Option<BorrowedFd<'sp>>) -> Self {
self.shared_sqpoll_fd = shared_sqpoll_fd;
self
}
pub fn build<'a, F>(
self,
buf_capacity: usize,
file_complete: F,
) -> io::Result<IoUringFileCreator<'a>>
where
F: FnMut(FileInfo) -> Option<File> + 'a,
{
let buf_capacity = buf_capacity.max(self.write_capacity as usize);
let buffer = PageAlignedMemory::new(buf_capacity)?;
self.build_with_buffer(buffer, file_complete)
}
fn build_with_buffer<'a, F: FnMut(FileInfo) -> Option<File> + 'a>(
self,
mut buffer: PageAlignedMemory,
file_complete: F,
) -> io::Result<IoUringFileCreator<'a>> {
let buf_capacity =
buffer.as_mut().len() / self.write_capacity as usize * self.write_capacity as usize;
assert_ne!(buf_capacity, 0, "write size aligned buffer is too small");
let buf_slice_mut = &mut buffer.as_mut()[..buf_capacity];
assert!(
self.write_capacity
.is_multiple_of(DIRECT_IO_WRITE_LEN_ALIGNMENT)
|| !self.write_with_direct_io,
"write capacity ({}) must be multiple of {DIRECT_IO_WRITE_LEN_ALIGNMENT} for direct IO",
self.write_capacity
);
let buffers = unsafe {
IoBufferChunk::split_buffer_chunks(
buf_slice_mut,
self.write_capacity,
self.register_buffer,
)
};
let state = FileCreatorState::new(buffers.collect(), file_complete);
let io_uring = self.create_io_uring(buf_capacity)?;
let ring = Ring::new(io_uring, state);
if self.register_buffer {
unsafe { IoBufferChunk::register(buf_slice_mut, &ring)? };
}
Ok(IoUringFileCreator {
ring,
write_with_direct_io: self.write_with_direct_io,
_backing_buffer: buffer,
})
}
fn create_io_uring(&self, buf_capacity: usize) -> io::Result<IoUring> {
let ring_qsize = self
.ring_squeue_size
.unwrap_or((buf_capacity / self.write_capacity as usize / 2).max(1) as u32);
let ring = sqpoll::io_uring_builder_with(self.shared_sqpoll_fd).build(ring_qsize)?;
ring.submitter()
.register_iowq_max_workers(&mut [self.max_iowq_workers, 1])?;
Ok(ring)
}
}
pub struct IoUringFileCreator<'a> {
ring: Ring<FileCreatorState<'a>, FileCreatorOp>,
write_with_direct_io: bool,
_backing_buffer: PageAlignedMemory,
}
impl FileCreator for IoUringFileCreator<'_> {
fn schedule_create_at_dir(
&mut self,
path: PathBuf,
mode: u32,
parent_dir_handle: Arc<File>,
contents: &mut dyn Read,
) -> io::Result<()> {
let file_key = self.open(path, mode, parent_dir_handle)?;
self.write_and_close(contents, file_key)
}
fn file_complete(&mut self, file: File, path: PathBuf, size: FileSize) {
let file_info = FileInfo { file, path, size };
(self.ring.context_mut().file_complete)(file_info);
}
fn drain(&mut self) -> io::Result<()> {
let res = self.ring.drain();
self.ring.context().log_stats();
res
}
}
impl IoUringFileCreator<'_> {
fn open(&mut self, path: PathBuf, mode: u32, dir_handle: Arc<File>) -> io::Result<usize> {
let file = PendingFile::from_path(path);
let path_cstring = Pin::new(file.path_cstring());
let file_key = self.wait_add_file(file)?;
let op = FileCreatorOp::Open(OpenOp {
dir_handle,
path_cstring,
mode,
file_key,
write_with_direct_io: self.write_with_direct_io,
});
self.ring.push(op)?;
Ok(file_key)
}
fn wait_add_file(&mut self, file: PendingFile) -> io::Result<usize> {
loop {
self.ring.process_completions()?;
if self.ring.context().files.len() < self.ring.context().files.capacity() {
break;
}
self.ring
.submit_and_wait(1, CHECK_PROGRESS_AFTER_SUBMIT_TIMEOUT)?;
}
let file_key = self.ring.context_mut().files.insert(file);
Ok(file_key)
}
fn write_and_close(&mut self, mut src: impl Read, file_key: usize) -> io::Result<()> {
let mut offset = 0;
let mut reached_eof = false;
while !reached_eof {
let buf = self.wait_free_buf()?;
let state = self.ring.context_mut();
let file_state = state.files.get_mut(file_key).unwrap();
let mut mut_slice =
unsafe { slice::from_raw_parts_mut(buf.as_mut_ptr(), buf.len() as usize) };
let mut write_len;
loop {
let len = src.read(mut_slice)?;
if len == 0 {
reached_eof = true;
write_len = buf.len() - mut_slice.len() as IoSize;
file_state.size_on_eof = Some(write_len as FileSize + offset);
break;
}
if len == mut_slice.len() {
write_len = buf.len();
break;
}
mut_slice = &mut mut_slice[len..];
}
if self.write_with_direct_io && write_len < buf.len() {
let align_truncated_write_len =
write_len / DIRECT_IO_WRITE_LEN_ALIGNMENT * DIRECT_IO_WRITE_LEN_ALIGNMENT;
if align_truncated_write_len != write_len {
file_state.non_dio_eof_write = Some(FinalNonDirectIoWrite {
buf: None,
file_offset: offset + align_truncated_write_len as FileSize,
buf_offset: align_truncated_write_len,
write_len: write_len - align_truncated_write_len,
});
write_len = align_truncated_write_len;
}
}
file_state.writes_started += 1;
if let Some(file) = &file_state.open_file {
let fd = types::Fd(file.as_raw_fd());
if write_len == 0 {
return FileCreatorState::mark_write_completed(
&mut self.ring,
file_key,
fd,
true,
buf,
);
}
let op = WriteOp {
file_key,
fd,
offset,
buf,
buf_offset: 0,
write_len,
};
state.submitted_writes_size += write_len as usize;
self.ring.push(FileCreatorOp::Write(op))?;
} else {
file_state.backlog.push((buf, offset, write_len));
}
offset += write_len as FileSize;
}
Ok(())
}
fn wait_free_buf(&mut self) -> io::Result<IoBufferChunk> {
loop {
self.ring.process_completions()?;
let state = self.ring.context_mut();
if let Some(buf) = state.buffers.pop_front() {
return Ok(buf);
}
state.stats.no_buf_count += 1;
state.stats.no_buf_sum_submitted_write_sizes += state.submitted_writes_size;
self.ring
.submit_and_wait(1, CHECK_PROGRESS_AFTER_SUBMIT_TIMEOUT)?;
}
}
}
struct FileCreatorState<'a> {
files: FixedSlab<PendingFile>,
buffers: VecDeque<IoBufferChunk>,
file_complete: Box<dyn FnMut(FileInfo) -> Option<File> + 'a>,
num_owned_files: usize,
submitted_writes_size: usize,
stats: FileCreatorStats,
}
impl<'a> FileCreatorState<'a> {
fn new(
buffers: VecDeque<IoBufferChunk>,
file_complete: impl FnMut(FileInfo) -> Option<File> + 'a,
) -> Self {
Self {
files: FixedSlab::with_capacity(MAX_OPEN_FILES),
buffers,
file_complete: Box::new(file_complete),
num_owned_files: 0,
submitted_writes_size: 0,
stats: FileCreatorStats::default(),
}
}
fn mark_file_opened(&mut self, file_key: usize, fd: types::Fd, direct_io: bool) -> BacklogVec {
let file = self.files.get_mut(file_key).unwrap();
file.open_file = Some(unsafe { File::from_raw_fd(fd.0) });
file.file_uses_direct_io = direct_io;
self.num_owned_files += 1;
if self.buffers.len() * 2 > self.buffers.capacity() {
self.stats.large_buf_headroom_count += 1;
}
mem::take(&mut file.backlog)
}
fn mark_write_completed(
ring: &mut impl RingAccess<Context = Self, Operation = FileCreatorOp>,
file_key: usize,
fd: types::Fd,
is_eof_write: bool,
buf: IoBufferChunk,
) -> io::Result<()> {
let this = ring.context_mut();
let file_state = this.files.get_mut(file_key).unwrap();
file_state.writes_completed += 1;
match file_state.non_dio_eof_write.as_mut() {
Some(eof_write) if is_eof_write => {
assert!(eof_write.buf.replace(buf).is_none())
}
_ => this.buffers.push_front(buf),
}
if file_state.required_writes_done() {
file_state.ensure_direct_io_disabled(fd)?;
if let Some(op) = file_state.try_take_final_write_op(file_key, fd) {
file_state.writes_started += 1;
this.submitted_writes_size += op.write_len as usize;
return ring.push(FileCreatorOp::Write(op));
}
if let Some(file_info) = file_state.take_completed_file_info() {
match (this.file_complete)(file_info) {
Some(unconsumed_file) => ring.push(FileCreatorOp::Close(CloseOp::new(
file_key,
unconsumed_file,
)))?,
None => this.mark_file_complete(file_key),
};
}
}
Ok(())
}
fn mark_file_complete(&mut self, file_key: usize) {
let _ = self.files.remove(file_key);
self.num_owned_files -= 1;
}
fn log_stats(&self) {
self.stats.log();
}
}
#[derive(Debug, Default)]
struct FileCreatorStats {
large_buf_headroom_count: u32,
no_buf_count: u32,
no_buf_sum_submitted_write_sizes: usize,
}
impl FileCreatorStats {
fn log(&self) {
let avg_writes_at_no_buf = self
.no_buf_sum_submitted_write_sizes
.checked_div(self.no_buf_count as usize)
.unwrap_or_default();
log::info!(
"files creation stats - large buf headroom: {}, no buf count: {}, avg pending writes \
at no buf: {avg_writes_at_no_buf}",
self.large_buf_headroom_count,
self.no_buf_count,
);
}
}
#[derive(Debug)]
struct OpenOp {
dir_handle: Arc<File>,
path_cstring: Pin<CString>,
mode: libc::mode_t,
file_key: usize,
write_with_direct_io: bool,
}
impl OpenOp {
fn entry(&mut self) -> squeue::Entry {
let at_dir_fd = types::Fd(self.dir_handle.as_raw_fd());
let mut flags = O_CREAT | O_TRUNC | O_NOFOLLOW | O_RDWR | DEFAULT_STATUS_FLAGS;
if self.write_with_direct_io {
flags |= O_DIRECT;
}
opcode::OpenAt::new(at_dir_fd, self.path_cstring.as_ptr() as _)
.flags(flags)
.mode(self.mode)
.build()
}
fn complete(
&mut self,
ring: &mut Completion<FileCreatorState, FileCreatorOp>,
res: io::Result<RawFd>,
) -> io::Result<()>
where
Self: Sized,
{
let fd = types::Fd(res?);
let backlog =
ring.context_mut()
.mark_file_opened(self.file_key, fd, self.write_with_direct_io);
for (buf, offset, write_len) in backlog {
if write_len == 0 {
return FileCreatorState::mark_write_completed(ring, self.file_key, fd, true, buf);
}
let op = WriteOp {
file_key: self.file_key,
fd,
offset,
buf,
buf_offset: 0,
write_len,
};
ring.context_mut().submitted_writes_size += write_len as usize;
ring.push(FileCreatorOp::Write(op))?;
}
Ok(())
}
}
#[derive(Debug)]
struct CloseOp {
file_key: usize,
fd: types::Fd,
}
impl<'a> CloseOp {
fn new(file_key: usize, file_to_close: File) -> Self {
let fd = types::Fd(file_to_close.into_raw_fd());
Self { file_key, fd }
}
fn entry(&mut self) -> squeue::Entry {
opcode::Close::new(self.fd).build()
}
fn complete(
&mut self,
ring: &mut Completion<FileCreatorState<'a>, FileCreatorOp>,
res: io::Result<i32>,
) -> io::Result<()>
where
Self: Sized,
{
let _ = res?;
ring.context_mut().mark_file_complete(self.file_key);
Ok(())
}
}
#[derive(Debug)]
struct WriteOp {
file_key: usize,
fd: types::Fd,
offset: FileSize,
buf: IoBufferChunk,
buf_offset: IoSize,
write_len: IoSize,
}
impl<'a> WriteOp {
fn entry(&mut self) -> squeue::Entry {
let WriteOp {
file_key: _,
fd,
offset,
buf,
buf_offset,
write_len,
} = self;
let buf_ptr = unsafe { buf.as_mut_ptr().byte_add(*buf_offset as usize) };
let entry = match buf.io_buf_index() {
Some(io_buf_index) => opcode::WriteFixed::new(*fd, buf_ptr, *write_len, io_buf_index)
.offset(*offset)
.ioprio(IO_PRIO_BE_HIGHEST)
.build(),
None => opcode::Write::new(*fd, buf_ptr, *write_len)
.offset(*offset)
.ioprio(IO_PRIO_BE_HIGHEST)
.build(),
};
entry.flags(squeue::Flags::ASYNC)
}
fn complete(
&mut self,
ring: &mut Completion<FileCreatorState<'a>, FileCreatorOp>,
res: io::Result<i32>,
) -> io::Result<()>
where
Self: Sized,
{
let written = match res {
Ok(0) => return Err(io::ErrorKind::WriteZero.into()),
Ok(res) => res as IoSize,
Err(err) => return Err(err),
};
let WriteOp {
file_key,
fd,
offset,
buf,
buf_offset,
write_len,
} = self;
let buf = mem::replace(buf, IoBufferChunk::empty());
ring.context_mut().submitted_writes_size -= written as usize;
let total_buf_written = *buf_offset + written;
if written < *write_len {
log::warn!("short write ({written}/{}), file={}", *write_len, *file_key);
return ring.push(FileCreatorOp::Write(WriteOp {
file_key: *file_key,
fd: *fd,
offset: *offset + written as FileSize,
buf,
buf_offset: total_buf_written,
write_len: *write_len - written,
}));
}
FileCreatorState::mark_write_completed(
ring,
*file_key,
*fd,
total_buf_written < buf.len(),
buf,
)
}
}
#[derive(Debug)]
enum FileCreatorOp {
Open(OpenOp),
Close(CloseOp),
Write(WriteOp),
}
impl RingOp<FileCreatorState<'_>> for FileCreatorOp {
fn entry(&mut self) -> squeue::Entry {
match self {
Self::Open(op) => op.entry(),
Self::Close(op) => op.entry(),
Self::Write(op) => op.entry(),
}
}
fn complete(
&mut self,
ring: &mut Completion<FileCreatorState, Self>,
res: io::Result<i32>,
) -> io::Result<()>
where
Self: Sized,
{
match self {
Self::Open(op) => op.complete(ring, res),
Self::Close(op) => op.complete(ring, res),
Self::Write(op) => op.complete(ring, res),
}
}
}
type PendingWrite = (IoBufferChunk, FileSize, IoSize);
#[derive(Debug)]
struct PendingFile {
path: PathBuf,
open_file: Option<File>,
backlog: BacklogVec,
size_on_eof: Option<FileSize>,
file_uses_direct_io: bool,
non_dio_eof_write: Option<FinalNonDirectIoWrite>,
writes_started: usize,
writes_completed: usize,
}
impl PendingFile {
fn from_path(path: PathBuf) -> Self {
Self {
path,
open_file: None,
backlog: SmallVec::new(),
writes_started: 0,
writes_completed: 0,
size_on_eof: None,
file_uses_direct_io: false,
non_dio_eof_write: None,
}
}
fn path_cstring(&self) -> CString {
let os_str = self.path.file_name().expect("path must contain filename");
CString::new(os_str.as_encoded_bytes()).expect("path mustn't contain interior NULs")
}
fn required_writes_done(&self) -> bool {
self.writes_started == self.writes_completed && self.source_fully_read()
}
fn source_fully_read(&self) -> bool {
self.size_on_eof.is_some()
}
fn ensure_direct_io_disabled(&mut self, fd: types::Fd) -> io::Result<()> {
if self.file_uses_direct_io {
let fcntl_res = unsafe { libc::fcntl(fd.0, libc::F_SETFL, DEFAULT_STATUS_FLAGS) };
if fcntl_res == -1 {
return Err(io::Error::last_os_error());
}
self.file_uses_direct_io = false;
}
Ok(())
}
fn try_take_final_write_op(&mut self, file_key: usize, fd: types::Fd) -> Option<WriteOp> {
let FinalNonDirectIoWrite {
buf,
file_offset: offset,
buf_offset,
write_len,
} = self.non_dio_eof_write.take()?;
let buf = buf.expect("should contain buffer for last write");
Some(WriteOp {
file_key,
fd,
offset,
buf,
buf_offset,
write_len,
})
}
fn take_completed_file_info(&mut self) -> Option<FileInfo> {
let size = self.size_on_eof.take().expect("content is fully read");
let file = self.open_file.take().expect("called once on completion");
let path = mem::take(&mut self.path);
Some(FileInfo { file, size, path })
}
}
#[derive(Debug)]
struct FinalNonDirectIoWrite {
buf: Option<IoBufferChunk>,
file_offset: FileSize,
buf_offset: IoSize,
write_len: IoSize,
}
#[cfg(test)]
mod tests {
use {super::*, std::io::Cursor, test_case::test_case};
#[test_case(0, 2 * 1024, 1024)]
#[test_case(1024, 1024, 1024)]
#[test_case(2 * 1024, 1024, 1024)]
#[test_case(4 * 1024, 2 * 1024, 1024)]
#[test_case(9 * 1024, 1024, 1024)]
#[test_case(9 * 1024, 2 * 1024, 1024)]
fn test_create_chunked_content(file_size: FileSize, buf_size: usize, write_size: IoSize) {
let contents = vec![1u8; file_size as usize];
let (contents_a, contents_b) = contents.split_at(file_size as usize / 3);
let mut contents = Cursor::new(contents_a).chain(contents_b);
let mut created = false;
let mut creator = IoUringFileCreatorBuilder::new()
.write_capacity(write_size)
.use_registered_buffers(false)
.build(buf_size, |file_info| {
assert_eq!(file_info.size, file_size);
assert_eq!(file_info.file.metadata().unwrap().len(), file_info.size);
created = true;
Some(file_info.file)
})
.unwrap();
let temp_dir = tempfile::tempdir().unwrap();
let dir = Arc::new(File::open(temp_dir.path()).unwrap());
let file_path = temp_dir.path().join("test.txt");
creator
.schedule_create_at_dir(file_path, 0o644, dir, &mut contents)
.unwrap();
creator.drain().unwrap();
drop(creator);
assert!(created);
}
#[test]
fn test_non_registered_buffer_create() {
let temp_dir = tempfile::tempdir().unwrap();
let mut read_data = vec![];
let callback = |mut fi: FileInfo| {
fi.file.read_to_end(&mut read_data).unwrap();
Some(fi.file)
};
let mut creator = IoUringFileCreatorBuilder::new()
.write_capacity(4 * 1024)
.use_registered_buffers(false)
.build(16 * 1024, callback)
.unwrap();
let dir = Arc::new(File::open(temp_dir.path()).unwrap());
let file_path = temp_dir.path().join("file.txt");
let file_size = 64 * 1024;
let data = (0..).take(file_size).map(|v| v as u8).collect::<Vec<_>>();
creator
.schedule_create_at_dir(file_path, 0o600, dir, &mut Cursor::new(&data))
.unwrap();
creator.drain().unwrap();
drop(creator);
assert_eq!(file_size, read_data.len());
assert_eq!(data, read_data);
}
#[test_case(4 * DIRECT_IO_WRITE_LEN_ALIGNMENT, 4 * DIRECT_IO_WRITE_LEN_ALIGNMENT)]
#[test_case(4 * DIRECT_IO_WRITE_LEN_ALIGNMENT, 2 * DIRECT_IO_WRITE_LEN_ALIGNMENT)]
#[test_case(5 * DIRECT_IO_WRITE_LEN_ALIGNMENT, 2 * DIRECT_IO_WRITE_LEN_ALIGNMENT)]
#[test_case(1000 + 4 * DIRECT_IO_WRITE_LEN_ALIGNMENT, 2 * DIRECT_IO_WRITE_LEN_ALIGNMENT)]
#[test_case(1000 + 5 * DIRECT_IO_WRITE_LEN_ALIGNMENT, 2 * DIRECT_IO_WRITE_LEN_ALIGNMENT)]
fn test_direct_io_create(file_size: IoSize, buf_size: IoSize) {
let mut read_data = Vec::with_capacity(file_size as usize);
let callback = |mut fi: FileInfo| {
fi.file.read_to_end(&mut read_data).unwrap();
Some(fi.file)
};
let mut creator = IoUringFileCreatorBuilder::new()
.write_capacity(8 * 1024)
.write_with_direct_io(true)
.build(buf_size as usize, callback)
.unwrap();
let temp_dir = tempfile::tempdir().unwrap();
let dir = Arc::new(File::open(temp_dir.path()).unwrap());
let file_path = temp_dir.path().join("file.txt");
let data: Vec<_> = (0..).take(file_size as usize).map(|v| v as u8).collect();
creator
.schedule_create_at_dir(file_path, 0o600, dir, &mut Cursor::new(&data))
.unwrap();
creator.drain().unwrap();
drop(creator);
assert_eq!(file_size, read_data.len() as IoSize);
assert_eq!(data, read_data);
}
}