io-uring: Create op_read helper function

Daksh14 · Daksh14 · commit b7ccee713971 · 2025-11-20T16:55:13.000+05:30
Fix tests fail on allocation failure

    Doc fix and fix double subtraction of offset

    Fix typos and use assert the stopped task's first poll
diff --git 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
diff --git a/tokio/src/fs/read_uring.rs b/tokio/src/fs/read_uring.rs
@@ -13,7 +13,7 @@ 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 blocks at time
+// 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>> {
@@ -27,141 +27,109 @@ 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, offset).await?;
-
-        if size_read == 0 {
-            return Ok(r_buf);
-        }
-
-        buf = r_buf;
-        fd = r_fd;
-        offset += size_read as u64;
-    }
-
     loop {
-        if buf.len() == buf.capacity() && buf.capacity() == start_cap {
+        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 (size_read, r_fd, r_buf) = small_probe_read(fd, buf, offset).await?;
+            let (size_read, r_fd, r_buf) = small_probe_read(fd, buf, &mut offset).await?;
 
             if size_read == 0 {
                 return Ok(r_buf);
             }
 
             buf = r_buf;
             fd = r_fd;
-            offset += size_read as u64;
         }
 
         // buf is full, need more capacity
         if buf.len() == buf.capacity() {
             buf.try_reserve(PROBE_SIZE)?;
         }
 
-        // doesn't matter if we have a valid size_hint or not, if we do more
-        // than 2 consecutive_short_reads, gradually increase the buffer
-        // capacity to read more data at a time
-
         // prepare the spare capacity to be read into
         let buf_len = usize::min(buf.spare_capacity_mut().len(), MAX_READ_SIZE);
 
-        // SAFETY: buf_len cannot be greater than u32::MAX because max_read_size
-        // is u32::MAX
-        let mut read_len = buf_len as u32;
-
-        loop {
-            // read into spare capacity
-            let (res, r_fd, r_buf) = Op::read(fd, buf, read_len, offset).await;
-
-            match res {
-                Ok(0) => return Ok(r_buf),
-                Ok(size_read) => {
-                    fd = r_fd;
-                    buf = r_buf;
-                    offset += size_read as u64;
-                    read_len -= size_read;
-
-                    // keep reading if there's something left to be read
-                    if read_len > 0 {
-                        continue;
-                    } else {
-                        break;
-                    }
-                }
-                Err(e) if e.kind() == ErrorKind::Interrupted => {
-                    buf = r_buf;
-                    fd = r_fd;
+        // buf_len cannot be greater than u32::MAX because MAX_READ_SIZE
+        // is less than u32::MAX
+        let read_len = buf_len as u32;
 
-                    continue;
-                }
-                Err(e) => return Err(e),
-            }
+        // read into spare capacity
+        let (size_read, r_fd, r_buf) = op_read(fd, buf, &mut offset, read_len).await?;
+
+        if size_read == 0 {
+            return Ok(r_buf);
         }
+
+        fd = r_fd;
+        buf = r_buf;
     }
 }
 
 async fn small_probe_read(
-    mut fd: OwnedFd,
+    fd: OwnedFd,
     mut buf: Vec<u8>,
-    offset: u64,
+    offset: &mut u64,
 ) -> io::Result<(u32, OwnedFd, Vec<u8>)> {
+    let read_len = PROBE_SIZE_U32;
+
     let mut temp_arr = [0; PROBE_SIZE];
-    let has_enough = buf.len() > PROBE_SIZE;
-
-    if has_enough {
-        // if we have more than PROBE_SIZE bytes in the buffer already then
-        // don't call reserve as we might potentially read 0 bytes
-        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);
-    } else {
-        // we don't even have PROBE_SIZE length in the buffer, we need this
-        // reservation
-        buf.reserve_exact(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 (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))
+}
+
+async fn op_read(
+    mut fd: OwnedFd,
+    mut buf: Vec<u8>,
+    offset: &mut u64,
+    mut read_len: u32,
+) -> io::Result<(u32, OwnedFd, Vec<u8>)> {
     loop {
-        let (res, r_fd, mut r_buf) = Op::read(fd, buf, PROBE_SIZE_U32, offset).await;
+        let (res, r_fd, r_buf) = Op::read(fd, buf, read_len, *offset).await;
 
         match res {
-            // return early if we inserted into reserved PROBE_SIZE
-            // bytes
-            Ok(size_read) if !has_enough => return Ok((size_read, r_fd, r_buf)),
+            Err(e) if e.kind() == ErrorKind::Interrupted => {
+                buf = r_buf;
+                fd = r_fd;
+            }
+            Err(e) => return Err(e),
             Ok(size_read) => {
-                let old_len = r_buf.len() - (size_read as usize);
+                *offset += size_read as u64;
+                read_len -= size_read;
 
-                r_buf.splice(old_len..old_len, temp_arr);
+                if read_len == 0 || size_read == 0 {
+                    return Ok((size_read, r_fd, r_buf));
+                }
 
-                return Ok((size_read, r_fd, r_buf));
-            }
-            Err(e) if e.kind() == ErrorKind::Interrupted => {
                 buf = r_buf;
                 fd = r_fd;
-
-                continue;
             }
-            Err(e) => return Err(e),
         }
     }
 }
diff --git a/tokio/src/io/uring/read.rs b/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.
diff --git a/tokio/tests/fs_uring_read.rs b/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();
 
