diff --git a/tokio/src/fs/mod.rs b/tokio/src/fs/mod.rs index d677645ec..8701d3a10 100644 --- a/tokio/src/fs/mod.rs +++ b/tokio/src/fs/mod.rs @@ -237,9 +237,6 @@ pub use self::metadata::metadata; mod open_options; pub use self::open_options::OpenOptions; -cfg_io_uring! { - pub(crate) use self::open_options::UringOpenOptions; -} mod read; pub use self::read::read; @@ -298,6 +295,13 @@ cfg_windows! { pub use self::symlink_file::symlink_file; } +cfg_io_uring! { + pub(crate) mod read_uring; + pub(crate) use self::read_uring::read_uring; + + pub(crate) use self::open_options::UringOpenOptions; +} + use std::io; #[cfg(not(test))] diff --git a/tokio/src/fs/read.rs b/tokio/src/fs/read.rs index 36edf2a5b..955c3592c 100644 --- a/tokio/src/fs/read.rs +++ b/tokio/src/fs/read.rs @@ -30,6 +30,16 @@ use std::{io, path::Path}; /// /// [`ErrorKind::Interrupted`]: std::io::ErrorKind::Interrupted /// +/// # io_uring support +/// +/// On Linux, you can also use io_uring for executing system calls. To enable +/// io_uring, you need to specify the `--cfg tokio_unstable` flag at compile time, +/// enable the io-uring cargo feature, and set the `Builder::enable_io_uring` +/// runtime option. +/// +/// Support for io_uring is currently experimental, so its behavior may change +/// or it may be removed in future versions. +/// /// # Examples /// /// ```no_run @@ -45,5 +55,23 @@ use std::{io, path::Path}; /// ``` pub async fn read(path: impl AsRef) -> io::Result> { let path = path.as_ref().to_owned(); + + #[cfg(all( + tokio_unstable, + feature = "io-uring", + feature = "rt", + feature = "fs", + target_os = "linux" + ))] + { + use crate::fs::read_uring; + + let handle = crate::runtime::Handle::current(); + let driver_handle = handle.inner.driver().io(); + if driver_handle.check_and_init()? { + return read_uring(&path).await; + } + } + asyncify(move || std::fs::read(path)).await } diff --git a/tokio/src/fs/read_uring.rs b/tokio/src/fs/read_uring.rs new file mode 100644 index 000000000..63eacd795 --- /dev/null +++ b/tokio/src/fs/read_uring.rs @@ -0,0 +1,134 @@ +use crate::fs::OpenOptions; +use crate::runtime::driver::op::Op; + +use std::io; +use std::io::ErrorKind; +use std::os::fd::OwnedFd; +use std::path::Path; + +// this algorithm is inspired from rust std lib version 1.90.0 +// https://doc.rust-lang.org/1.90.0/src/std/io/mod.rs.html#409 +const PROBE_SIZE: usize = 32; +const PROBE_SIZE_U32: u32 = PROBE_SIZE as u32; + +// Max bytes we can read using io uring submission at a time +// SAFETY: cannot be higher than u32::MAX for safe cast +// Set to read max 64 MiB at time +const MAX_READ_SIZE: usize = 64 * 1024 * 1024; + +pub(crate) async fn read_uring(path: &Path) -> io::Result> { + let file = OpenOptions::new().read(true).open(path).await?; + + // TODO: use io uring in the future to obtain metadata + let size_hint: Option = file.metadata().await.map(|m| m.len() as usize).ok(); + + let fd: OwnedFd = file + .try_into_std() + .expect("unexpected in-flight operation detected") + .into(); + + let mut buf = Vec::new(); + + if let Some(size_hint) = size_hint { + buf.try_reserve(size_hint)?; + } + + read_to_end_uring(fd, buf).await +} + +async fn read_to_end_uring(mut fd: OwnedFd, mut buf: Vec) -> io::Result> { + let mut offset = 0; + let start_cap = buf.capacity(); + + loop { + if buf.len() == buf.capacity() && buf.capacity() == start_cap && buf.len() >= PROBE_SIZE { + // The buffer might be an exact fit. Let's read into a probe buffer + // and see if it returns `Ok(0)`. If so, we've avoided an + // unnecessary increasing of the capacity. But if not, append the + // probe buffer to the primary buffer and let its capacity grow. + let (r_fd, r_buf, is_eof) = small_probe_read(fd, buf, &mut offset).await?; + + if is_eof { + return Ok(r_buf); + } + + buf = r_buf; + fd = r_fd; + } + + // buf is full, need more capacity + if buf.len() == buf.capacity() { + buf.try_reserve(PROBE_SIZE)?; + } + + // prepare the spare capacity to be read into + let buf_len = usize::min(buf.spare_capacity_mut().len(), MAX_READ_SIZE); + + // buf_len cannot be greater than u32::MAX because MAX_READ_SIZE + // is less than u32::MAX + let read_len = u32::try_from(buf_len).expect("buf_len must always fit in u32"); + + // read into spare capacity + let (r_fd, r_buf, is_eof) = op_read(fd, buf, &mut offset, read_len).await?; + + if is_eof { + return Ok(r_buf); + } + + fd = r_fd; + buf = r_buf; + } +} + +async fn small_probe_read( + fd: OwnedFd, + mut buf: Vec, + offset: &mut u64, +) -> io::Result<(OwnedFd, Vec, bool)> { + let read_len = PROBE_SIZE_U32; + + let mut temp_arr = [0; PROBE_SIZE]; + // we don't call this function if buffer's length < PROBE_SIZE + let back_bytes_len = buf.len() - PROBE_SIZE; + + temp_arr.copy_from_slice(&buf[back_bytes_len..]); + + // We're decreasing the length of the buffer and len is greater + // than PROBE_SIZE. So we can read into the discarded length + buf.truncate(back_bytes_len); + + let (r_fd, mut r_buf, is_eof) = op_read(fd, buf, offset, read_len).await?; + // If `size_read` returns zero due to reasons such as buffer's exact fit, + // then this `try_reserve` does not perform allocation. + r_buf.try_reserve(PROBE_SIZE)?; + r_buf.splice(back_bytes_len..back_bytes_len, temp_arr); + + Ok((r_fd, r_buf, is_eof)) +} + +// Takes a amount of length to read and returns a singluar read in the buffer +// +// Returns the file descriptor, buffer and EOF reached or not +async fn op_read( + mut fd: OwnedFd, + mut buf: Vec, + offset: &mut u64, + read_len: u32, +) -> io::Result<(OwnedFd, Vec, bool)> { + loop { + let (res, r_fd, r_buf) = Op::read(fd, buf, read_len, *offset).await; + + match res { + Err(e) if e.kind() == ErrorKind::Interrupted => { + buf = r_buf; + fd = r_fd; + } + Err(e) => return Err(e), + Ok(size_read) => { + *offset += size_read as u64; + + return Ok((r_fd, r_buf, size_read == 0)); + } + } + } +} diff --git a/tokio/src/io/uring/mod.rs b/tokio/src/io/uring/mod.rs index 4899d0a4a..facad596f 100644 --- a/tokio/src/io/uring/mod.rs +++ b/tokio/src/io/uring/mod.rs @@ -1,3 +1,4 @@ pub(crate) mod open; +pub(crate) mod read; pub(crate) mod utils; pub(crate) mod write; diff --git a/tokio/src/io/uring/read.rs b/tokio/src/io/uring/read.rs new file mode 100644 index 000000000..e8ee633ac --- /dev/null +++ b/tokio/src/io/uring/read.rs @@ -0,0 +1,61 @@ +use crate::runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op}; + +use io_uring::{opcode, types}; +use std::io::{self, Error}; +use std::os::fd::{AsRawFd, OwnedFd}; + +#[derive(Debug)] +pub(crate) struct Read { + fd: OwnedFd, + buf: Vec, +} + +impl Completable for Read { + type Output = (io::Result, OwnedFd, Vec); + + fn complete(self, cqe: CqeResult) -> Self::Output { + let mut buf = self.buf; + + if let Ok(len) = cqe.result { + let new_len = buf.len() + len as usize; + // SAFETY: Kernel read len bytes + unsafe { buf.set_len(new_len) }; + } + + (cqe.result, self.fd, buf) + } + + fn complete_with_error(self, err: Error) -> Self::Output { + (Err(err), self.fd, self.buf) + } +} + +impl Cancellable for Read { + fn cancel(self) -> CancelData { + CancelData::Read(self) + } +} + +impl Op { + // Submit a request to read a FD at given length and offset into a + // dynamic buffer with uninitialized memory. The read happens on uninitialized + // buffer and no overwriting happens. + + // SAFETY: The `len` of the amount to be read and the buffer that is passed + // should have capacity > len. + // + // If `len` read is higher than vector capacity then setting its length by + // the caller in terms of size_read can be unsound. + pub(crate) fn read(fd: OwnedFd, mut buf: Vec, len: u32, offset: u64) -> Self { + // don't overwrite on already written part + assert!(buf.spare_capacity_mut().len() >= len as usize); + let buf_mut_ptr = buf.spare_capacity_mut().as_mut_ptr().cast(); + + let read_op = opcode::Read::new(types::Fd(fd.as_raw_fd()), buf_mut_ptr, len) + .offset(offset) + .build(); + + // SAFETY: Parameters are valid for the entire duration of the operation + unsafe { Op::new(read_op, Read { fd, buf }) } + } +} diff --git a/tokio/src/runtime/driver/op.rs b/tokio/src/runtime/driver/op.rs index 37945cf5a..d2b9289ce 100644 --- a/tokio/src/runtime/driver/op.rs +++ b/tokio/src/runtime/driver/op.rs @@ -1,4 +1,5 @@ use crate::io::uring::open::Open; +use crate::io::uring::read::Read; use crate::io::uring::write::Write; use crate::runtime::Handle; @@ -17,6 +18,7 @@ use std::task::{Context, Poll, Waker}; pub(crate) enum CancelData { Open(Open), Write(Write), + Read(Read), } #[derive(Debug)] diff --git a/tokio/tests/fs_uring_read.rs b/tokio/tests/fs_uring_read.rs new file mode 100644 index 000000000..4ce6d454a --- /dev/null +++ b/tokio/tests/fs_uring_read.rs @@ -0,0 +1,202 @@ +//! Uring file operations tests. + +#![cfg(all( + tokio_unstable, + feature = "io-uring", + feature = "rt", + feature = "fs", + target_os = "linux" +))] + +use futures::future::Future; +use std::future::poll_fn; +use std::io::Write; +use std::path::PathBuf; +use std::sync::mpsc; +use std::task::{Context, Poll, Waker}; +use std::time::Duration; +use tempfile::NamedTempFile; +use tokio::fs::read; +use tokio::runtime::{Builder, Runtime}; +use tokio_test::assert_pending; +use tokio_util::task::TaskTracker; + +fn multi_rt(n: usize) -> Box Runtime> { + Box::new(move || { + Builder::new_multi_thread() + .worker_threads(n) + .enable_all() + .build() + .unwrap() + }) +} + +fn current_rt() -> Box Runtime> { + Box::new(|| Builder::new_current_thread().enable_all().build().unwrap()) +} + +fn rt_combinations() -> Vec Runtime>> { + vec![ + current_rt(), + multi_rt(1), + multi_rt(2), + multi_rt(8), + multi_rt(64), + multi_rt(256), + ] +} + +#[test] +fn shutdown_runtime_while_performing_io_uring_ops() { + fn run(rt: Runtime) { + let (done_tx, done_rx) = mpsc::channel(); + let (_tmp, path) = create_tmp_files(1); + // keep 100 permits + const N: i32 = 100; + rt.spawn(async move { + let path = path[0].clone(); + + // spawning a bunch of uring operations. + let mut futs = vec![]; + + // spawning a bunch of uring operations. + for _ in 0..N { + let path = path.clone(); + let mut fut = Box::pin(read(path)); + + poll_fn(|cx| { + assert_pending!(fut.as_mut().poll(cx)); + Poll::<()>::Pending + }) + .await; + + futs.push(fut); + } + + tokio::task::yield_now().await; + }); + + std::thread::spawn(move || { + rt.shutdown_timeout(Duration::from_millis(300)); + done_tx.send(()).unwrap(); + }); + + done_rx.recv().unwrap(); + } + + for rt in rt_combinations() { + run(rt()); + } +} + +#[test] +fn read_many_files() { + fn run(rt: Runtime) { + const NUM_FILES: usize = 512; + + let (_tmp_files, paths): (Vec, Vec) = create_tmp_files(NUM_FILES); + + rt.block_on(async move { + let tracker = TaskTracker::new(); + + for i in 0..10_000 { + let path = paths.get(i % NUM_FILES).unwrap().clone(); + tracker.spawn(async move { + let bytes = read(path).await.unwrap(); + assert_eq!(bytes, vec![20; 1023]); + }); + } + tracker.close(); + tracker.wait().await; + }); + } + + for rt in rt_combinations() { + run(rt()); + } +} + +#[tokio::test] +async fn read_small_large_files() { + let (_tmp, path) = create_large_temp_file(); + + let bytes = read(path).await.unwrap(); + + assert_eq!(bytes, create_buf(5000)); + + let (_tmp, path) = create_small_temp_file(); + + let bytes = read(path).await.unwrap(); + + assert_eq!(bytes, create_buf(20)); +} + +#[tokio::test] +async fn cancel_op_future() { + let (_tmp_file, path): (Vec, Vec) = create_tmp_files(1); + let path = path[0].clone(); + + let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel(); + + let handle = tokio::spawn(async move { + let fut = read(path.clone()); + tokio::pin!(fut); + + poll_fn(move |_| { + // If io_uring is enabled (and not falling back to the thread pool), + // the first poll should return Pending. + assert_pending!(fut.as_mut().poll(&mut Context::from_waker(Waker::noop()))); + tx.send(true).unwrap(); + + Poll::<()>::Pending + }) + .await; + }); + + // Wait for the first poll + + let val = rx.recv().await; + assert!(val.unwrap()); + + handle.abort(); + + let res = handle.await.unwrap_err(); + assert!(res.is_cancelled()); +} + +fn create_tmp_files(num_files: usize) -> (Vec, Vec) { + let mut files = Vec::with_capacity(num_files); + for _ in 0..num_files { + let mut tmp = NamedTempFile::new().unwrap(); + let buf = vec![20; 1023]; + tmp.write_all(&buf).unwrap(); + let path = tmp.path().to_path_buf(); + files.push((tmp, path)); + } + + files.into_iter().unzip() +} + +fn create_large_temp_file() -> (NamedTempFile, PathBuf) { + let mut tmp = NamedTempFile::new().unwrap(); + let buf = create_buf(5000); + + tmp.write_all(&buf).unwrap(); + let path = tmp.path().to_path_buf(); + + (tmp, path) +} + +fn create_small_temp_file() -> (NamedTempFile, PathBuf) { + let mut tmp = NamedTempFile::new().unwrap(); + let buf = create_buf(20); + + tmp.write_all(&buf).unwrap(); + let path = tmp.path().to_path_buf(); + + (tmp, path) +} + +fn create_buf(length: usize) -> Vec { + (0..length).map(|i| i as u8).collect() +}