From 4a47719219f3f3f5a4aa48cb57d2d39184d98c40 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Mon, 31 Aug 2026 23:37:56 -0400 Subject: [PATCH 1/4] introduce drop guarde for shrink_to_fit --- arrow-buffer/src/buffer/immutable.rs | 88 +++++++++++++++++++++++++--- 1 file changed, 79 insertions(+), 9 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 12d873ba2c91..b1ac0bb700fe 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -224,17 +224,33 @@ 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) } + // Drop guard: keeps Buffer::ptr consistent with Bytes::ptr even if a custom + // MemoryReservation::resize panics inside try_realloc (see #10379). + struct PtrSync<'a> { + ptr: &'a mut *const u8, + bytes: *const Bytes, + offset: usize, + is_empty: bool, + } + impl Drop for PtrSync<'_> { + fn drop(&mut self) { + // SAFETY: bytes is valid while Arc is held by Buffer + let base = unsafe { (*self.bytes).as_ptr() }; + *self.ptr = if self.is_empty { + base + } else { + // SAFETY: offset is within the allocated region + unsafe { base.add(self.offset) } + }; } - } else { - // Failure to reallocate is fine; we just failed to free up memory. } + let _sync = PtrSync { + ptr: &mut self.ptr, + bytes: bytes as *const Bytes, + offset, + is_empty, + }; + bytes.try_realloc(desired_capacity).ok(); } } @@ -1158,4 +1174,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()); + } } From 03409e14c3f670b161aacc952ff730b37bb43d0f Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Mon, 31 Aug 2026 23:58:35 -0400 Subject: [PATCH 2/4] use better variable names --- arrow-buffer/src/buffer/immutable.rs | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index b1ac0bb700fe..f0db1e2bb33a 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -227,16 +227,16 @@ impl Buffer { // Drop guard: keeps Buffer::ptr consistent with Bytes::ptr even if a custom // MemoryReservation::resize panics inside try_realloc (see #10379). struct PtrSync<'a> { - ptr: &'a mut *const u8, - bytes: *const Bytes, + buffer_ptr: &'a mut *const u8, + bytes_ptr: *const Bytes, offset: usize, is_empty: bool, } impl Drop for PtrSync<'_> { fn drop(&mut self) { - // SAFETY: bytes is valid while Arc is held by Buffer - let base = unsafe { (*self.bytes).as_ptr() }; - *self.ptr = if self.is_empty { + // SAFETY: bytes_ptr is valid while Arc is held by Buffer + let base = unsafe { (*self.bytes_ptr).as_ptr() }; + *self.buffer_ptr = if self.is_empty { base } else { // SAFETY: offset is within the allocated region @@ -245,8 +245,8 @@ impl Buffer { } } let _sync = PtrSync { - ptr: &mut self.ptr, - bytes: bytes as *const Bytes, + buffer_ptr: &mut self.ptr, + bytes_ptr: std::ptr::from_ref::(bytes), offset, is_empty, }; From 9cc95a6a51525b342e2a68a7247159501d6f8989 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Tue, 1 Sep 2026 23:40:53 -0400 Subject: [PATCH 3/4] aviod double mut accessor --- arrow-buffer/src/buffer/immutable.rs | 33 +++++++--------------------- arrow-buffer/src/bytes.rs | 11 +++++++++- 2 files changed, 18 insertions(+), 26 deletions(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index f0db1e2bb33a..1ba4a0a33f66 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -224,33 +224,16 @@ impl Buffer { if desired_capacity < self.capacity() && let Some(bytes) = Arc::get_mut(&mut self.data) { - // Drop guard: keeps Buffer::ptr consistent with Bytes::ptr even if a custom - // MemoryReservation::resize panics inside try_realloc (see #10379). - struct PtrSync<'a> { - buffer_ptr: &'a mut *const u8, - bytes_ptr: *const Bytes, - offset: usize, - is_empty: bool, - } - impl Drop for PtrSync<'_> { - fn drop(&mut self) { - // SAFETY: bytes_ptr is valid while Arc is held by Buffer - let base = unsafe { (*self.bytes_ptr).as_ptr() }; - *self.buffer_ptr = if self.is_empty { - base + bytes + .try_realloc(desired_capacity, |base| { + self.ptr = if is_empty { + base.as_ptr() } else { - // SAFETY: offset is within the allocated region - unsafe { base.add(self.offset) } + // SAFETY: we kept all elements leading up to the offset + unsafe { base.as_ptr().add(offset) } }; - } - } - let _sync = PtrSync { - buffer_ptr: &mut self.ptr, - bytes_ptr: std::ptr::from_ref::(bytes), - offset, - is_empty, - }; - bytes.try_realloc(desired_capacity).ok(); + }) + .ok(); } } 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")] { From 08aaa63d672efed047b59b7fde4a06c31f9de8c5 Mon Sep 17 00:00:00 2001 From: RIchard Baah Date: Wed, 2 Sep 2026 08:39:04 -0400 Subject: [PATCH 4/4] re-add comment --- arrow-buffer/src/buffer/immutable.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-buffer/src/buffer/immutable.rs b/arrow-buffer/src/buffer/immutable.rs index 1ba4a0a33f66..4faf60856963 100644 --- a/arrow-buffer/src/buffer/immutable.rs +++ b/arrow-buffer/src/buffer/immutable.rs @@ -233,7 +233,7 @@ impl Buffer { unsafe { base.as_ptr().add(offset) } }; }) - .ok(); + .ok(); // Failure to reallocate is fine; we just failed to free up memory. } }