Fix tests fail on allocation failure

Doc fix and fix double subtraction of offset

Fix shutdown op test

Fix typos and use assert the stopped task's first poll

Author: Daksh14

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

tokio/src/fs/read_uring.rs

@@ -27,32 +27,19 @@ pub(crate) async fn read_uring(path: &Path) -> io::Result<Vec<u8>> {
         .expect("unexpected in-flight operation detected")
         .into();
 
-    // extra single capacity for the whole size to fit without any reallocation
-    let buf = Vec::with_capacity(size_hint.unwrap_or(0));
+    let mut buf = Vec::new();
 
-    read_to_end_uring(size_hint, fd, buf).await
+    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(
-    size_hint: Option<usize>,
-    mut fd: OwnedFd,
-    mut buf: Vec<u8>,
-) -> io::Result<Vec<u8>> {
+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();
 
-    // if buffer has no room and no size_hint, start with a small probe_read from 0 offset
-    if (size_hint.is_none() || size_hint == Some(0)) && buf.capacity() - buf.len() < PROBE_SIZE {
-        let (size_read, r_fd, r_buf) = small_probe_read(fd, buf, &mut offset).await?;
-
-        if size_read == 0 {
-            return Ok(r_buf);
-        }
-
-        fd = r_fd;
-        buf = r_buf;
-    }
-
     loop {
         if buf.len() == buf.capacity() && buf.capacity() == start_cap {
             // The buffer might be an exact fit. Let's read into a probe buffer
@@ -67,7 +54,6 @@ async fn read_to_end_uring(
 
             buf = r_buf;
             fd = r_fd;
-            offset += size_read as u64;
         }
 
         // buf is full, need more capacity
@@ -79,11 +65,11 @@ async fn read_to_end_uring(
         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 u32::MAX
-        let mut read_len = buf_len as u32;
+        // is less than u32::MAX
+        let read_len = buf_len as u32;
 
         // read into spare capacity
-        let (size_read, r_fd, r_buf) = op_read(fd, buf, &mut offset, &mut read_len).await?;
+        let (size_read, r_fd, r_buf) = op_read(fd, buf, &mut offset, read_len).await?;
 
         if size_read == 0 {
             return Ok(r_buf);
@@ -99,12 +85,12 @@ async fn small_probe_read(
     mut buf: Vec<u8>,
     offset: &mut u64,
 ) -> io::Result<(u32, OwnedFd, Vec<u8>)> {
-    let mut read_len = PROBE_SIZE_U32;
+    let read_len = PROBE_SIZE_U32;
 
     if buf.len() < PROBE_SIZE {
         buf.try_reserve(PROBE_SIZE)?;
 
-        return op_read(fd, buf, offset, &mut read_len).await;
+        return op_read(fd, buf, offset, read_len).await;
     }
 
     let mut temp_arr = [0; PROBE_SIZE];
@@ -116,8 +102,9 @@ async fn small_probe_read(
     // than PROBE_SIZE. So we can read into the discarded length
     buf.truncate(back_bytes_len);
 
-    let (size_read, r_fd, mut r_buf) = op_read(fd, buf, offset, &mut read_len).await?;
+    let (size_read, r_fd, mut r_buf) = op_read(fd, buf, offset, read_len).await?;
 
+    r_buf.try_reserve(PROBE_SIZE)?;
     r_buf.splice(back_bytes_len..back_bytes_len, temp_arr);
 
     Ok((size_read, r_fd, r_buf))
@@ -127,10 +114,10 @@ async fn op_read(
     mut fd: OwnedFd,
     mut buf: Vec<u8>,
     offset: &mut u64,
-    read_len: &mut u32,
+    mut read_len: u32,
 ) -> io::Result<(u32, OwnedFd, Vec<u8>)> {
     loop {
-        let (res, r_fd, r_buf) = Op::read(fd, buf, *read_len, *offset).await;
+        let (res, r_fd, r_buf) = Op::read(fd, buf, read_len, *offset).await;
 
         match res {
             Err(e) if e.kind() == ErrorKind::Interrupted => {
@@ -140,9 +127,9 @@ async fn op_read(
             Err(e) => return Err(e),
             Ok(size_read) => {
                 *offset += size_read as u64;
-                *read_len -= size_read;
+                read_len -= size_read;
 
-                if *read_len == 0 || size_read == 0 {
+                if read_len == 0 || size_read == 0 {
                     return Ok((size_read, r_fd, r_buf));
                 }
 

tokio/src/io/uring/read.rs

@@ -38,8 +38,8 @@ impl Cancellable for Read {
 
 impl Op<Read> {
     // Submit a request to read a FD at given length and offset into a
-    // dynamic buffer with uinitialized memory. The read happens on unitialized
-    // buffer and no overwiting happens.
+    // dynamic buffer with uninitialized memory. The read happens on unitialized
+    // buffer and no overwriting happens.
 
     // SAFETY: The `len` of the amount to be read and the buffer that is passed
     // should have capacity > len.

tokio/tests/fs_uring_read.rs

@@ -8,17 +8,17 @@
     target_os = "linux"
 ))]
 
-use futures::future::FutureExt;
+use futures::future::Future;
+use std::future::poll_fn;
 use std::io::Write;
+use std::path::PathBuf;
 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::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> {
@@ -49,34 +49,38 @@ fn rt_combinations() -> Vec<Box<dyn Fn() -> Runtime>> {
 #[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);
+        // keep 100 permits
+        const N: i32 = 100;
         rt.spawn(async move {
             let path = path[0].clone();
 
             // spawning a bunch of uring operations.
-            loop {
+            let mut futs = vec![];
+
+            // spawning a bunch of uring operations.
+            for _ in 0..N {
                 let path = path.clone();
-                tokio::spawn(async move {
-                    let bytes = read(path).await.unwrap();
+                let mut fut = Box::pin(read(path));
 
-                    assert_eq!(bytes, vec![20; 1023]);
-                });
+                poll_fn(|cx| {
+                    assert_pending!(fut.as_mut().poll(cx));
+                    Poll::<()>::Pending
+                })
+                .await;
 
-                // Avoid busy looping.
-                tokio::task::yield_now().await;
+                futs.push(fut);
             }
+
+            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();
     }
 
@@ -130,25 +134,29 @@ async fn read_small_large_files() {
 #[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 {
-        poll_fn(|cx| {
-            let fut = read(&path[0]);
+        let fut = read(path.clone());
+        tokio::pin!(fut);
 
+        poll_fn(move |cx| {
             // 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();
+            assert_pending!(fut.as_mut().poll(cx));
+            tx.send(true).unwrap();
 
             Poll::<()>::Pending
         })
         .await;
     });
 
     // Wait for the first poll
-    rx.recv().await.unwrap();
+
+    let val = rx.recv().await;
+    assert!(val.unwrap());
 
     handle.abort();