io-uring: Create op_read helper function

Daksh14 · Daksh14 · commit 7d4ccebb3f76 · 2025-11-29T15:50:29.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

    Add comments and panic on failed u32 conversion
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,111 @@ 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 = u32::try_from(buf_len).expect("buf_len must always fit in 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);
-    }
+    // 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 (size_read, r_fd, mut r_buf) = 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((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();
 
