io-uring: Implement tokio::fs::read

Author: Daksh14

tokio/src/fs/read.rs

@@ -45,5 +45,58 @@ 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"
+    ))]
+    {
+        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
 }
+
+#[cfg(all(
+    tokio_unstable,
+    feature = "io-uring",
+    feature = "rt",
+    feature = "fs",
+    target_os = "linux"
+))]
+async fn read_uring(path: &Path) -> io::Result<Vec<u8>> {
+    use crate::{fs::OpenOptions, runtime::driver::op::Op};
+    use std::os::fd::OwnedFd;
+
+    let file = OpenOptions::new().read(true).open(path).await?;
+
+    let size = file
+        .metadata()
+        .await
+        .map(|m| m.len() as usize)
+        .ok()
+        .unwrap_or(0);
+
+    let buf = Vec::with_capacity(size);
+
+    let fd: OwnedFd = file
+        .try_into_std()
+        .expect("unexpected in-flight operation detected")
+        .into();
+
+    let (read_size, mut buf) = Op::read(fd, buf)?.await?;
+
+    // SAFETY:
+    // 1. The buffer is initialized with `size` capacity
+    // 2. the read_size is the number of bytes read from the file
+    unsafe { buf.set_len(read_size as _) };
+
+    Ok(buf)
+}

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,42 @@
+use crate::runtime::driver::op::{CancelData, Cancellable, Completable, CqeResult, Op};
+use io_uring::{opcode, types};
+use std::{
+    io,
+    os::fd::{AsRawFd, OwnedFd},
+};
+
+#[derive(Debug)]
+pub(crate) struct Read {
+    buf: Vec<u8>,
+    fd: OwnedFd,
+}
+
+impl Completable for Read {
+    type Output = (u32, Vec<u8>);
+    fn complete(self, cqe: CqeResult) -> io::Result<Self::Output> {
+        let res = cqe.result?;
+
+        Ok((res, self.buf))
+    }
+}
+
+impl Cancellable for Read {
+    fn cancel(self) -> CancelData {
+        CancelData::Read(self)
+    }
+}
+
+impl Op<Read> {
+    /// Submit a request to open a file.
+    pub(crate) fn read(fd: OwnedFd, mut buf: Vec<u8>) -> io::Result<Self> {
+        let buf_mut_ptr = buf.as_mut_ptr();
+        let cap = buf.capacity() as _;
+
+        let read_op = opcode::Read::new(types::Fd(fd.as_raw_fd()), buf_mut_ptr, cap).build();
+
+        // SAFETY: Parameters are valid for the entire duration of the operation
+        let op = unsafe { Op::new(read_op, Read { fd, buf }) };
+
+        Ok(op)
+    }
+}

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,155 @@
+//! Uring file operations tests.
+
+#![cfg(all(
+    tokio_unstable,
+    feature = "io-uring",
+    feature = "rt",
+    feature = "fs",
+    target_os = "linux"
+))]
+
+use futures::future::FutureExt;
+use std::io::Write;
+use std::sync::mpsc;
+use std::task::Poll;
+use std::time::Duration;
+use std::{future::poll_fn, path::PathBuf};
+use tempfile::NamedTempFile;
+use tokio::{
+    fs::read,
+    runtime::{Builder, Runtime},
+};
+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 (tx, rx) = mpsc::channel();
+        let (done_tx, done_rx) = mpsc::channel();
+
+        let (_tmp, path) = create_tmp_files(1);
+        rt.spawn(async move {
+            let path = path[0].clone();
+
+            // spawning a bunch of uring operations.
+            loop {
+                let path = path.clone();
+                tokio::spawn(async move {
+                    let bytes = read(path).await.unwrap();
+
+                    assert_eq!(bytes, vec![20; 2]);
+                });
+
+                // Avoid busy looping.
+                tokio::task::yield_now().await;
+            }
+        });
+
+        std::thread::spawn(move || {
+            let rt: Runtime = rx.recv().unwrap();
+            rt.shutdown_timeout(Duration::from_millis(300));
+            done_tx.send(()).unwrap();
+        });
+
+        tx.send(rt).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; 2]);
+                });
+            }
+            tracker.close();
+            tracker.wait().await;
+        });
+    }
+
+    for rt in rt_combinations() {
+        run(rt());
+    }
+}
+
+#[tokio::test]
+async fn cancel_op_future() {
+    let (_tmp_file, path): (Vec<NamedTempFile>, Vec<PathBuf>) = create_tmp_files(1);
+
+    let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
+    let handle = tokio::spawn(async move {
+        poll_fn(|cx| {
+            let fut = read(&path[0]);
+
+            // If io_uring is enabled (and not falling back to the thread pool),
+            // the first poll should return Pending.
+            let _pending = Box::pin(fut).poll_unpin(cx);
+
+            tx.send(()).unwrap();
+
+            Poll::<()>::Pending
+        })
+        .await;
+    });
+
+    // Wait for the first poll
+    rx.recv().await.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; 2];
+        tmp.write_all(&buf).unwrap();
+        let path = tmp.path().to_path_buf();
+        files.push((tmp, path));
+    }
+
+    files.into_iter().unzip()
+}