Fix perf regression in Read::read_to_end on short reads due to not checking if the cursor has initialized bytes

[asder8215]

View all comments

This PR fixes #158008.

In particular, in #150129, it refactored some code within library/std/io/mod.rs to utilize BorrowedBuf::is_init instead of manually checking read_buf.init_len() == buf_len to see if the read buffer had initialized bytes. However, the BorrowedBuf is never marked or set as init within this function, and I think this portion of the code:

 // SAFETY: These bytes were initialized but not filled in the previous loop
unsafe {
     read_buf.set_init(initialized);
}

was removed by mistake. This PR reverts the changes made by #150129, so that we can mark the BorrowedBuf/read_buf as initialized using BorrowedBuf::set_init if in a previous iteration the cursor has initialized bytes. This would allow max_read_size to not be marked as usize::max if the read buffer contains initialized bytes.

[rustbot]

r? @Darksonn

rustbot has assigned @Darksonn.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

Why was this reviewer chosen?

The reviewer was selected based on:

• Owners of files modified in this PR: @ChrisDenton, libs
• @ChrisDenton, libs expanded to 12 candidates
• Random selection from Darksonn, Mark-Simulacrum, clarfonthey, jhpratt

[Darksonn]

This comment was added in #150129. Should it be removed?

View changes since the review

[asder8215]

Unsure. Are these comments still relevant @a1phyr?

[a1phyr]

Well, probably not if you start tracking initialized bytes :)

[Darksonn]

@rustbot author

[rustbot]

Reminder, once the PR becomes ready for a review, use @rustbot ready.

[a1phyr]

I didn't track uninitialized bytes in my previous MR because I thought it would be to complicated to do properly.

For example, if this MR improve some existing cases, it will not really solve the pathological case you sent in your issue for larger sizes (eg around a million): when you initialized N bytes, on the next round you will have N-1 spare initialized bytes left, so you won't be able to use set_init() (or you could initialize the rest manually).

All in all, it was a trade-off between code complexity, properly handling common cases but having suboptimal (but still acceptable) behavior in weird cases.

View changes since this review

[a1phyr]

This is true only if read_buf.is_init() (use the boolean below to avoid lifetime issues)

[a1phyr]

This condition is wrong: you compare the old buffer capacity and the new buffer spare capacity, but the start of the buffer has changed since then. It would be less error prone to track initialized bytes counting from the beginning of the Vec.

[Darksonn]

Yeah, please store the full capacity of the vector, including anything already written.

[Darksonn]

@rustbot author

[asder8215]

@rustbot ready

[Darksonn]

This is still using the wrong condition. Consider this scenario:

1. We set initialized_len = 10 - 5 where capacity=10, length=5.
2. The vector is resized and the capacity is now 15.
3. We read another 5 bytes with is_init=false.
4. The next iteration compares initialized_len with 15-10 which is true, and considers the buffer initialized.

But in this scenario the buffer was resized and not initialized, so this is wrong. It's really important that initialized_len is such that after the vector is resized, the comparison can never return true until you get a read with is_init=true.

Instead of storing capacity() - len() please just store capacity(). That way, if the capacity changes, then the capacity will no longer be the same.

@rustbot author

[asder8215]

Instead of storing capacity() - len() please just store capacity(). That way, if the capacity changes, then the capacity will no longer be the same.

Just to clarify, resizing only occurs potentially in spots where we call small_probe_reads or in the conditional where buf.len() == buf.capacity() and it does a try_reserve? It should still be safe to use buf_unfilled_len = capacity() - len() on the portion of code where we're converting the spare portion of the buffer into a BorrowedBuf/BorrowedCursor as it shouldn't be able to resize this borrowed buffer/slice with read_buf?

[Darksonn]

It's not obvious to me that that's the case. Maybe you are right, but I can't easily follow the logic.

To me, it would be a lot easier to figure out that the code is correct if you stored the capacity. If the vector was reallocated, then I know that the capacity changed, and therefore I know that we will not call set_init().

[asder8215]

@rustbot ready

[Darksonn]

I think it would be less confusing to also set is_init to false here.

View changes since the review

[Darksonn]

This variable keeps track of what operations we have done in the past, but I think it's generally easier to think about variables that keep track of facts about the world.

So for example, could we store the length of the subset of the buffer that we are currently reading into? Before each iteration you do if buf_len > spare.len() { buf_len = spare.len(); }. And after each iteration you do buf_len -= bytes_read followed by if buf_len == 0 { buf_len = max_read_size; }. The is_init variable then keeps track of whether the current subset is initialized or not.

View changes since the review

[asder8215]

I ended up refactoring and using the is_init variable and spare buffer length to determine whether we need to initialize more bytes into the spare buffer or not.

If our spare buffer has a remainder from the modulo operation with max_read_size, it should mean that there are bytes there that we have initialized; otherwise, if we see a 0, then that indicates that we need to initialize more bytes into the spare buffer (aside from the case when our spare buffer is the minimum between itself and the max_read_size in which case taking the mod of that with itself would always produce 0 and anyways we just need to present that remaining spare buffer length).

[rustbot]

This PR was rebased onto a different main commit. Here's a range-diff highlighting what actually changed.

Rebasing is a normal part of keeping PRs up to date, so no action is needed—this note is just to help reviewers.

[Darksonn]

I believe you have to set is_init to false here too? Since this may increase the number of bytes passed in the next iteration.

View changes since the review

[asder8215]

Doesn't is_init get set to false in the next loop? I believe this conditional doubles max_read_size when it notices that we filled the spare buffer with max_read_size bytes in one iteration rather than through multiple iterations. From the next loop, I think the buf.len() == buf.capacity() conditional earlier should be triggered, resize the buffer, and set is_init to false.

Also, now that I realize this, I think it would be impossible for buf_len > max_read_size because we already accounted for re-using any remaining uninitialized space in the spare buffer rather than creating more uninitialized bytes. I could change this conditional to just bytes_read == max_read_size?

[asder8215]

This conditional right here (which is effectively buf.len() == start_cap due to transitivity) only seems to be ran once to check if the buffer's length is the same as the starting buffer capacity (the ideal case here is that small_probe_read returns 0 and hence we can say for sure that whatever we read into the buffer was an exact fit + avoided doubling the buffer's capacity). Otherwise, if there were more bytes to read, then our buffer's capacity grows and it should be impossible for buf.len() == start_cap (start_cap does not change); therefore, I took this out of the loop since we don't need to run through this multiple times.

View changes since the review

[asder8215]

Still need was_init to set BorrowedBuf init field to true, and I use it down below so that we're not repeatedly assigning init_until with the same value unless we reach a point where we need to initialize more bytes into the spare buffer.

View changes since the review

[Darksonn]

Here's a thought: if init_until - buf.len() is positive but very small, we probably want to increase the buffer size even if was_init becomes false. Previously PROBE_SIZE was used a threshold for this. How about doing this?

let buf_len = if init_until > buf.len() + PROBE_SIZE {
    init_until - buf.len()
} else {
    usize::min(max_read_size, buf.capacity() - buf.len())
};
let was_init = init_until >= buf.len() + buf_len;

[asder8215]

Kept the exact fit check inside the buf.len() == buf.capacity() conditional, so that we're not running through both conditional separately/repeatedly in the case that buf.len() != buf.capacity(). Hope that's fine.

View changes since the review

[Darksonn]

I think the check has to be done twice. When small_probe_read calls extend_from_slice, that will reallocate and then push more bytes. Since more bytes are pushed after the reallocation, there is a possible scenario where the reallocation increases the size by the same amount read, which leads to buf.len() == buf.capacity() being the case after the call to small_probe_read.

Therefore, you must check buf.len() == buf.capacity() again after calling small_probe_read. Otherwise you will call read_buf on an empty slice below, which doesn't work.

[asder8215]

Right and in that case you'll end up having bytes_read = 0, which will cause an early return even though there may still be more bytes to read.

[ChrisDenton]

I think this might have some overlap with #142872 (that PR is a year old but I was just reminded of it)