fs: support io_uring with tokio::fs::read (#7696)

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))]

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<Path>) -> io::Result<Vec<u8>> {
     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
 }

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<Vec<u8>> {
+    let file = OpenOptions::new().read(true).open(path).await?;
+
+    // TODO: use io uring in the future to obtain metadata
+    let size_hint: Option<usize> = 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<u8>) -> io::Result<Vec<u8>> {
+    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<u8>,
+    offset: &mut u64,
+) -> io::Result<(OwnedFd, Vec<u8>, 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<u8>,
+    offset: &mut u64,
+    read_len: u32,
+) -> io::Result<(OwnedFd, Vec<u8>, 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));
+            }
+        }
+    }
+}

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;

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<u8>,
+}
+
+impl Completable for Read {
+    type Output = (io::Result<u32>, OwnedFd, Vec<u8>);
+
+    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<Read> {
+    // 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<u8>, 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 }) }
+    }
+}

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)]

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<dyn Fn() -> Runtime> {
+    Box::new(move || {
+        Builder::new_multi_thread()
+            .worker_threads(n)
+            .enable_all()
+            .build()
+            .unwrap()
+    })
+}
+
+fn current_rt() -> Box<dyn Fn() -> Runtime> {
+    Box::new(|| Builder::new_current_thread().enable_all().build().unwrap())
+}
+
+fn rt_combinations() -> Vec<Box<dyn Fn() -> 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<NamedTempFile>, Vec<PathBuf>) = 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<NamedTempFile>, Vec<PathBuf>) = 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<NamedTempFile>, Vec<PathBuf>) {
+    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<u8> {
+    (0..length).map(|i| i as u8).collect()
+}