fix(arrow-buffer): add a fallible collect_bool - #10984
Conversation
collect_bool reserves ceil(len / 64) u64 words before calling the closure, so a len taken from a row or bit count in user controlled data aborts the process rather than failing. The reservation used Vec::with_capacity, which has no way to report that. Add try_collect_bool returning Result<Self, MutableBufferError> and implement collect_bool on top of it, matching the with_capacity and try_with_capacity pairing already in this file. The reservation now goes through Vec::try_reserve and maps a failure onto the existing AllocationError variant.
| buf.reserve(usize::MAX); | ||
| } | ||
| #[test] | ||
| fn try_collect_bool_reports_a_len_it_cannot_reserve() { |
There was a problem hiding this comment.
@Cintu07 this is causing the miri to fail
test buffer::mutable::tests::try_collect_bool_reports_a_len_it_cannot_reserve ...
stderr ───
warning: extern crate `criterion` is unused in crate `arrow_buffer`
--> arrow-buffer/src/lib.rs:18:1
|
18 | //! Low-level buffer abstractions for [Apache Arrow Rust](https://docs.rs/arrow)
| ^
|
= help: remove the dependency or add `use criterion as _;` to the crate root
= note: requested on the command line with `--force-warn unused-crate-dependencies`
error: resource exhaustion: tried to allocate more memory than available to compiler
--> /home/runner/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/alloc/src/raw_vec/mod.rs:568:13
|
568 | self.alloc.allocate(new_layout)
| ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ resource exhaustion occurred here
|
There was a problem hiding this comment.
thanks, fixed in a2e0417. #[cfg_attr(miri, ignore)] on that one the way the rest of
the repo does it. the other test only uses small sizes so it should still run under
miri, ci will confirm.
There was a problem hiding this comment.
I'd prefer to not skip the miri check here.
Im not sure that we need either of these test 🤔. I think the fix here is self explanatory enough to not warrant them.
Miri interprets the allocation rather than handing it to the system allocator, so reserving 2^60 bits exhausts the interpreter itself instead of coming back as an error. Ignore that test under miri, the way the other oversized and inline-assembly tests in the repo already are. The small size test keeps running there, so the new code path is still covered under miri.
| buf.reserve(usize::MAX); | ||
| } | ||
| #[test] | ||
| fn try_collect_bool_reports_a_len_it_cannot_reserve() { |
There was a problem hiding this comment.
I'd prefer to not skip the miri check here.
Im not sure that we need either of these test 🤔. I think the fix here is self explanatory enough to not warrant them.
| } | ||
|
|
||
| #[test] | ||
| fn try_collect_bool_matches_collect_bool_for_sizes_that_fit() { |
There was a problem hiding this comment.
nit: I don't think try_collect_bool does enough thats different from collect_bool to warrant adding another test
|
|
||
| /// Fallible version of [`MutableBuffer::collect_bool`]. | ||
| /// | ||
| /// `len` is a bit count, so the reservation is `ceil(len / 64)` words. Reserving that |
There was a problem hiding this comment.
nit: This is true if “words” means u64 storage words, but it’s a bit ambiguous on 32-bit targets where “word” usually means 32 bits. Maybe say ceil(len / 64) u64 slots instead.
There was a problem hiding this comment.
done in b58729e, reads ceil(len / 64) u64 slots now.
Per review: the two tests are removed rather than skipped under miri, and the reservation doc says u64 slots instead of words, which was ambiguous on 32 bit targets.
Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com>
Co-authored-by: Jeffrey Vo <jeffrey.vo.australia@gmail.com>
Which issue does this PR close?
MutableBuffer::collect_boolallocates fromlenbefore any iteration #10973.Rationale for this change
collect_boolreserves ceil(len / 64) u64 words before it calls the closure once, so alenthat comes from a row or bit count in user controlled data aborts the process rather than failing. the reservation went throughVec::with_capacity, which has no way to report that.this file already has the shape for it.
with_capacityistry_with_capacity(..).unwrap_or_else(|e| panic!("{e}"))at line 154, and there are seven other fallible twins beside it:try_from_len_zeroed,try_reserve,try_repeat_slice_n_times,try_resize,try_shrink_to_fit,try_extend_from_sliceandtry_extend_zeros.collect_boolis the allocating entry point that does not have one.What changes are included in this PR?
try_collect_boolreturningResult<Self, MutableBufferError>collect_boolcalls it and unwraps, the same waywith_capacitycallstry_with_capacityVec::try_reserve, and a failure maps onto the existingAllocationErrorvariant carrying the layout it tried to takeno behaviour change for any
lenthat already worked.collect_boolstill panics, one level down, and its doc comment now says so.i went with a fallible entry point rather than capping the reservation, since a cap only moves where the abort happens rather than letting a caller handle it. happy to do the cap instead if you would rather not grow the api surface.
Are these changes tested?
yes, two tests.
try_collect_bool_reports_a_len_it_cannot_reserveuses the 2^60 from the issue and asserts it comes back asAllocationErrorand that the closure ran zero times. the reproducer in the issue aborts on that same input withmemory allocation of 144115188075855872 bytes failed.try_collect_bool_matches_collect_bool_for_sizes_that_fitchecks the two agree byte for byte at 0, 1, 63, 64, 65 and 1000 bits, since 64 is the word boundary and 8 is the truncation boundary.arrow-bufferis 351 passed,arrow-arrayis 721 passed, and fmt and clippy with-D warningsare clean. i ran arrow-array as well becausecollect_boolhas around twenty call sites in it and this changes how the function is built.Are there any user-facing changes?
try_collect_boolis new and additive.collect_boolkeeps its signature and still panics on alenit cannot reserve.