io-uring: Create op_read helper function

Author: Daksh14

tokio/src/fs/read_uring.rs

@@ -39,20 +39,18 @@ async fn read_to_end_uring(
     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?;
+        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 = r_buf;
     }
 
     loop {
@@ -61,7 +59,7 @@ async fn read_to_end_uring(
             // 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);
@@ -88,43 +86,26 @@ async fn read_to_end_uring(
         // 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;
+        // read into spare capacity
+        let (size_read, r_fd, r_buf) = op_read(fd, buf, &mut offset, &mut read_len).await?;
 
-                    continue;
-                }
-                Err(e) => return Err(e),
-            }
+        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 mut temp_arr = [0; PROBE_SIZE];
+    let mut read_len = PROBE_SIZE_U32;
+
     let has_enough = buf.len() > PROBE_SIZE;
 
     if has_enough {
@@ -141,27 +122,47 @@ async fn small_probe_read(
         buf.reserve_exact(PROBE_SIZE);
     }
 
+    let (size_read, r_fd, mut r_buf) = op_read(fd, buf, offset, &mut read_len).await?;
+
+    // return early if we inserted into reserved PROBE_SIZE
+    // bytes
+    if !has_enough {
+        Ok((size_read, r_fd, r_buf))
+    } else {
+        let old_len = r_buf.len() - (size_read as usize);
+
+        r_buf.splice(old_len..old_len, temp_arr);
+
+        Ok((size_read, r_fd, r_buf))
+    }
+}
+
+async fn op_read(
+    mut fd: OwnedFd,
+    mut buf: Vec<u8>,
+    offset: &mut u64,
+    read_len: &mut 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),
         }
     }
 }