diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 12d873ba2c91..4faf60856963 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -224,17 +224,16 @@ impl Buffer { if desired_capacity < self.capacity() && let Some(bytes) = Arc::get_mut(&mut self.data) { - if bytes.try_realloc(desired_capacity).is_ok() { - // Realloc complete - update our pointer into `bytes`: - self.ptr = if is_empty { - bytes.as_ptr() - } else { - // SAFETY: we kept all elements leading up to the offset - unsafe { bytes.as_ptr().add(offset) } - } - } else { - // Failure to reallocate is fine; we just failed to free up memory. - } + bytes + .try_realloc(desired_capacity, |base| { + self.ptr = if is_empty { + base.as_ptr() + } else { + // SAFETY: we kept all elements leading up to the offset + unsafe { base.as_ptr().add(offset) } + }; + }) + .ok(); // Failure to reallocate is fine; we just failed to free up memory. } } @@ -1158,4 +1157,58 @@ mod tests { assert_eq!(buffer_back.as_slice(), expected.as_slice()); } } + + #[test] + #[cfg(feature = "pool")] + fn test_shrink_to_fit_panicking_reservation() { + use std::panic::{AssertUnwindSafe, catch_unwind}; + + use crate::pool::{MemoryPool, MemoryReservation}; + + #[derive(Debug)] + struct PanicPool; + + #[derive(Debug)] + struct PanicReservation { + panicked: bool, + } + + impl MemoryReservation for PanicReservation { + fn size(&self) -> usize { + 0 + } + fn resize(&mut self, _: usize) { + if !self.panicked { + self.panicked = true; + panic!("intentional panic in resize"); + } + } + } + + impl MemoryPool for PanicPool { + fn reserve(&self, _: usize) -> Box { + Box::new(PanicReservation { panicked: false }) + } + fn available(&self) -> isize { + isize::MAX + } + fn used(&self) -> usize { + 0 + } + fn capacity(&self) -> usize { + usize::MAX + } + } + + let pool = PanicPool; + let data: Vec = (0..8).collect(); + let mut buf = Buffer::from_slice_ref(data.as_slice()); + buf.claim(&pool); + + // shrink_to_fit panics because PanicReservation::resize panics, but + // Buffer::ptr must stay consistent with Bytes::ptr (no use-after-free). + let _ = catch_unwind(AssertUnwindSafe(|| buf.shrink_to_fit())); + + assert_eq!(buf.as_slice(), data.as_slice()); + } } diff --git a/arrow-buffer/src/bytes.rs b/arrow-buffer/src/bytes.rs index de9f7befe6e8..deadb5f48f65 100644 --- a/arrow-buffer/src/bytes.rs +++ b/arrow-buffer/src/bytes.rs @@ -134,7 +134,15 @@ impl Bytes { /// Returns `Err` if the memory was allocated with a custom allocator, /// or the call to `realloc` failed, for whatever reason. /// In case of `Err`, the [`Bytes`] will remain as it was (i.e. have the old size). - pub(crate) fn try_realloc(&mut self, new_len: usize) -> Result<(), ()> { + /// + /// `on_reallocated` is called after [`Bytes`] has updated its internal + /// pointer, but before resizing the memory reservation, which may call user + /// code. + pub(crate) fn try_realloc( + &mut self, + new_len: usize, + on_reallocated: impl FnOnce(NonNull), + ) -> Result<(), ()> { if let Deallocation::Standard(old_layout) = self.deallocation { if old_layout.size() == new_len { return Ok(()); // Nothing to do @@ -162,6 +170,7 @@ impl Bytes { self.ptr = ptr; self.len = new_len; self.deallocation = Deallocation::Standard(new_layout); + on_reallocated(ptr); #[cfg(feature = "pool")] {