fix: bound IPC message allocations by the bytes actually read - #10522
fix: bound IPC message allocations by the bytes actually read#10522ranflarion wants to merge 2 commits into
Conversation
| /// and the resulting allocation failure aborts the process instead of surfacing an error | ||
| /// the caller can handle. | ||
| #[test] | ||
| fn test_stream_reader_rejects_implausible_body_length() { |
There was a problem hiding this comment.
Please add a test that exercises a body with length shorter than the message.bodyLength, to test the early exit on EOF path
There was a problem hiding this comment.
added, test_stream_reader_rejects_truncated_body covers a 1024-byte declared body backed by 10 bytes, which fails before the first growth step, and a MAX_PREALLOC_BYTES + 1 body backed by exactly MAX_PREALLOC_BYTES, which fails after one, so the resize path is exercised before hitting EOF.
|
run benchmark ipc_reader |
This comment was marked as duplicate.
This comment was marked as duplicate.
|
🤖 Arrow criterion benchmark completed (GKE) | trigger Instance: Comparing bound-ipc-body-allocation (c1a91eb) to ed92960 (merge-base) diff Run configurationrun benchmark ipc_readerCPU Details (lscpu)Details
Resource Usagebase (merge-base)
branch
File an issue against this benchmark runner |
|
| const MAX_PREALLOC_BYTES: usize = 64 * 1024 * 1024; | ||
|
|
||
| /// Reads exactly `len` bytes of message body, without reserving `len` before reading it. | ||
| fn read_body_bounded<R: Read>(reader: &mut R, len: usize) -> Result<MutableBuffer, ArrowError> { |
There was a problem hiding this comment.
this seems reasonable -- basically it reallocates in 64MB chunks when reading
for a 1GB message, this will result in quite a few reallocations, but I suppose the idea is that such message sizes are rare.
I was thinking about some usecase reading 8K rows record batch where each row has a 1MB document -- that is 8GB and would likely reallocate 128K times, which is probably not great for performance
Is there some way to use a built in rust function (or maybe employ a doubling strategy) or something?
There was a problem hiding this comment.
MutableBuffer::resize goes through reserve, which already grows capacity to max(required, capacity * 2) (arrow-buffer/src/buffer/mutable.rs:256), so the 64MB steps only moved len while capacity doubled underneath: the 8GB case reallocates 7 times (64M→128M→…→8G) with cumulative copy under one extra pass over the body, not 128K times. I've made the doubling explicit in read_body_bounded anyway (resize(len.min(target * 2))), so the growth policy is visible in the loop instead of relying on reserve internals, and the loop runs log2(n / 64MB) iterations. Total zeroing is unchanged from before this PR, from_len_zeroed zeroed the whole body up front too. On a built-in, take(len).read_to_end is what the metadata path uses, but read_to_end only fills a Vec and the body has to stay a MutableBuffer for the 64-byte aligned allocation, so the body loop stays hand-rolled.
yes i missed that, #9869 is after the same three changes (take/read_to_end for the metadata, try_from on bodyLength, bounded body reads). The difference is the body path: #9869 copies through a 64KiB scratch chunk via extend_from_slice, so every body pays a second full copy, while here the bytes land directly in the destination MutableBuffer. |
5 participants
Which issue does this PR close?
Rationale for this change
MessageReader::maybe_nextreserves both the metadata length and the message body length before reading any of the bytes they describe. Both come out of the stream itself, so a corrupted or truncated stream is handed straight to the allocator:MutableBuffer::from_len_zeroed(message.bodyLength() as usize)on an implausible length either aborts the process (memory allocation of N bytes failed, which is not catchable and takes the host process with it) or panics onLayoutError. A negativebodyLengthis accepted too, sinceas usizewraps it to a large positive length.I hit this fuzzing real IPC blocks rather than crafted ones: single-bit flips over the framing region of genuine streams produced
memory allocation of 1125899907497992 bytes failedandSIGABRT. Where those blocks cross disk or a network, one flipped bit ends the process instead of failing a read the caller could retry.What changes are included in this PR?
bodyLengthnow goes throughusize::try_from, so a negative length is a parse error rather than a huge positive one.Neither length reserves more than
MAX_PREALLOC_BYTES(64 MiB) before the bytes behind it have arrived. Bodies up to that size are allocated in one go exactly as before; larger ones grow as the data arrives, which costs the reallocations thatMutableBuffer::reservedoubling implies. That constant is the one judgement call here, trading the size of the bounded allocation a malformed stream can still ask for against how large a body keeps the single-allocation path, so it is worth a second opinion.The metadata read switches from
resize(meta_len, 0)plusread_exacttotake(meta_len).read_to_end(&mut self.buf). That reuses the retained capacity across messages and drops the zeroing entirely, so it should be slightly cheaper than what it replaces rather than a cost, andTakereturnsOk(0)at its limit so there is no extra read.Only the streaming path is touched.
read_blockon the file side has the same shape atarrow-ipc/src/reader.rs:875and twounwrap()s on block metadata besides; I left it alone to keep this reviewable, and noted it in the issue.This overlaps #9777, which is after the same zeroing for performance reasons. The two want the same thing here, and I am happy to rebase onto whatever lands first.
Are these changes tested?
Yes, two tests in
arrow-ipc/src/reader.rs.test_stream_reader_rejects_implausible_body_lengthcoversi64::MAX,1 << 50and-1;test_stream_reader_rejects_unbacked_metadata_lengthcovers a metadata length ofi32::MAXwith nineteen bytes behind it. Both fail without the change: the first panics insideMutableBuffer::from_len_zeroed, and the second spends 7.7s zeroing 2 GiB before reporting the wrong error.The existing
arrow-ipcsuite passes (139 tests), along withcargo fmt --all --checkandcargo clippy -p arrow-ipc --all-targets --all-features -- -D warnings.Are there any user-facing changes?
No API changes. A stream that previously aborted or panicked now returns an
ArrowError. No breaking changes.