Repository: tokio-rs/tokio
File: tokio/src/fs/read_uring.rs (Rust)

Code before:
use crate::fs::OpenOptions;
use crate::runtime::driver::op::Op;

use std::io::ErrorKind;
use std::os::fd::OwnedFd;
use std::path::Path;
use std::{cmp, io};

// this algorithm is inspired from rust std lib version 1.90.0
// https://doc.rust-lang.org/1.90.0/src/std/io/mod.rs.html#409
const PROBE_SIZE: usize = 32;
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
const MAX_READ_SIZE: usize = u32::MAX as usize;

pub(crate) async fn read_uring(path: &Path) -> io::Result<Vec<u8>> {
    let file = OpenOptions::new().read(true).open(path).await?;

    // TODO: use io uring in the future to obtain metadata
    let size_hint: Option<usize> = file
        .metadata()
        .await
        .map(|m| usize::try_from(m.len()).unwrap_or(usize::MAX))
        .ok();

    let fd: OwnedFd = file
        .try_into_std()
        .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).saturating_add(1));

    read_to_end_uring(size_hint, fd, buf).await
}

async fn read_to_end_uring(
    size_hint: Option<usize>,
    mut fd: OwnedFd,
    mut buf: Vec<u8>,
) -> io::Result<Vec<u8>> {
    let mut offset = 0;
    let mut consecutive_short_reads = 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?;

Review comment (performance):
The purpose of `size_hint` is to improve performance by avoiding memory allocation. However, it involves thread pool mechanisms—which may not be more efficient than multiple memory allocations due to the overhead of creating/terminating threads and inter-thread synchronization operations.

We need performance data to demonstrate that this approach is faster than the method without using `size_hint`.

I asked this because I want to eliminate unnecessary logic as much as possible.

Code after:
use crate::fs::OpenOptions;
use crate::runtime::driver::op::Op;

use std::io;
use std::io::ErrorKind;
use std::os::fd::OwnedFd;
use std::path::Path;

// this algorithm is inspired from rust std lib version 1.90.0
// https://doc.rust-lang.org/1.90.0/src/std/io/mod.rs.html#409
const PROBE_SIZE: usize = 32;
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 MiB at time
const MAX_READ_SIZE: usize = 64 * 1024 * 1024;

pub(crate) async fn read_uring(path: &Path) -> io::Result<Vec<u8>> {
    let file = OpenOptions::new().read(true).open(path).await?;

    // TODO: use io uring in the future to obtain metadata
    let size_hint: Option<usize> = file.metadata().await.map(|m| m.len() as usize).ok();

    let fd: OwnedFd = file
        .try_into_std()
        .expect("unexpected in-flight operation detected")
        .into();

    let mut buf = Vec::new();

    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(mut fd: OwnedFd, mut buf: Vec<u8>) -> io::Result<Vec<u8>> {
    let mut offset = 0;
    let start_cap = buf.capacity();

    loop {
        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 (r_fd, r_buf, is_eof) = small_probe_read(fd, buf, &mut offset).await?;

            if is_eof {