Repository: tokio-rs/tokio
File: tokio/src/io/util/read_to_end.rs (Rust)

Code before:
    // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every
    // time is 4,500 times (!) slower than this if the reader has a very small
    // amount of data to return.
    const NUM_BYTES: usize = 32;
    let try_small_read = buf.try_small_read_first(NUM_BYTES);

    // Get a ReadBuf into the vector.
    let mut read_buf = buf.get_read_buf();
    let filled_before = read_buf.filled().len();
    let poll_result;

    let filled_after = if try_small_read {
        let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] =
            unsafe { MaybeUninit::uninit().assume_init() };
        let mut small_read_buf = ReadBuf::uninit(&mut small_buf);
        poll_result = read.poll_read(cx, &mut small_read_buf);
        let filled = small_read_buf.filled();
        read_buf.put_slice(filled);
        filled_before + filled.len()
    } else {
        buf.reserve(NUM_BYTES);
        read_buf = buf.get_read_buf();
        poll_result = read.poll_read(cx, &mut read_buf);
        read_buf.filled().len()
    };

    // Update the length of the vector using the result of poll_read.
    let read_buf_parts = into_read_buf_parts(read_buf);
    buf.apply_read_buf(read_buf_parts);
    let n = filled_after - filled_before;

    match poll_result {
        Poll::Pending => {
            // In this case, nothing should have been read. However we still
            // update the vector in case the poll_read call initialized parts of
            // the vector's unused capacity.
            debug_assert_eq!(filled_before, filled_after);
            Poll::Pending
        }
        Poll::Ready(Err(err)) => {
            debug_assert_eq!(filled_before, filled_after);
            Poll::Ready(Err(err))
        }
        Poll::Ready(Ok(())) => Poll::Ready(Ok(n)),
    }
}

impl<A> Future for ReadToEnd<'_, A>
where
    A: AsyncRead + ?Sized + Unpin,
{

Review comment (bug):
You can't use `put_slice` here. It panics if the vector doesn't have enough space to hold the bytes.

Please add a test that catches this bug.

Code after:
    // of data to return. Simply tacking on an extra DEFAULT_BUF_SIZE space every
    // time is 4,500 times (!) slower than this if the reader has a very small
    // amount of data to return. When the vector is full with its starting
    // capacity, we first try to read into a small buffer to see if we reached
    // an EOF. This only happens when the starting capacity is >= NUM_BYTES, since
    // we allocate at least NUM_BYTES each time. This avoids the unnecessary
    // allocation that we attempt before reading into the vector.

    const NUM_BYTES: usize = 32;
    let try_small_read = buf.try_small_read_first(NUM_BYTES);

    // Get a ReadBuf into the vector.
    let mut read_buf;
    let poll_result;

    let n = if try_small_read {
        // Read some bytes using a small read.
        let mut small_buf: [MaybeUninit<u8>; NUM_BYTES] = [MaybeUninit::uninit(); NUM_BYTES];
        let mut small_read_buf = ReadBuf::uninit(&mut small_buf);
        poll_result = read.poll_read(cx, &mut small_read_buf);
        let to_write = small_read_buf.filled();

        // Ensure we have enough space to fill our vector with what we read.
        read_buf = buf.get_read_buf();
        if to_write.len() > read_buf.remaining() {
            buf.reserve(NUM_BYTES);
            read_buf = buf.get_read_buf();
        }
        read_buf.put_slice(to_write);

        to_write.len()
    } else {
        // Ensure we have enough space for reading.
        buf.reserve(NUM_BYTES);
        read_buf = buf.get_read_buf();

        // Read data directly into vector.
        let filled_before = read_buf.filled().len();
        poll_result = read.poll_read(cx, &mut read_buf);

        // Compute the number of bytes read.
        read_buf.filled().len() - filled_before
    };

    // Update the length of the vector using the result of poll_read.
    let read_buf_parts = into_read_buf_parts(read_buf);
    buf.apply_read_buf(read_buf_parts);

    match poll_result {
        Poll::Pending => {
            // In this case, nothing should have been read. However we still