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

Code before:

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?;

        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 {
            // 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?;

            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() {
            if consecutive_short_reads > 1 {
                buf.try_reserve(PROBE_SIZE.saturating_mul(consecutive_short_reads))?;
            } else {
                buf.try_reserve(PROBE_SIZE)?;
            }
        }

        // doesn't matter if we have a valid size_hint or not, if we do more

Review comment (refactor):
I think we should just not call `small_probe_read` if the length is less than `SMALL_PROBE_READ`. In that case, we can just call `buf.reserve(SMALL_PROBE_READ)` and continue to the normal code below. This will simplify `small_probe_read` because `has_enough` will always be true so you can remove those codepaths.

Code after:
    }

    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 {
                return Ok(r_buf);
            }

            buf = r_buf;
            fd = r_fd;
        }

        // buf is full, need more capacity
        if buf.len() == buf.capacity() {
            buf.try_reserve(PROBE_SIZE)?;
        }

        // prepare the spare capacity to be read into
        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 less than u32::MAX
        let read_len = u32::try_from(buf_len).expect("buf_len must always fit in u32");

        // read into spare capacity
        let (r_fd, r_buf, is_eof) = op_read(fd, buf, &mut offset, read_len).await?;

        if is_eof {
            return Ok(r_buf);
        }

        fd = r_fd;
        buf = r_buf;
    }
}

async fn small_probe_read(
    fd: OwnedFd,