From 4f290aa562465674dceef8516da06ec79383448f Mon Sep 17 00:00:00 2001 From: cetra3 Date: Wed, 8 Jul 2026 11:44:58 +0930 Subject: [PATCH 1/2] Fix pool claim race condition --- arrow-buffer/src/buffer/mutable.rs | 49 +++------- arrow-buffer/src/bytes.rs | 22 ++--- arrow-buffer/src/pool.rs | 147 ++++++++++++++++++++++++++--- 3 files changed, 155 insertions(+), 63 deletions(-) diff --git a/arrow-buffer/src/buffer/mutable.rs b/arrow-buffer/src/buffer/mutable.rs index b6e6a70c6cba..b6f2c34f6d69 100644 --- a/arrow-buffer/src/buffer/mutable.rs +++ b/arrow-buffer/src/buffer/mutable.rs @@ -27,9 +27,7 @@ use crate::{ }; #[cfg(feature = "pool")] -use crate::pool::{MemoryPool, MemoryReservation}; -#[cfg(feature = "pool")] -use std::sync::Mutex; +use crate::pool::{MemoryPool, TrackedReservation}; use super::Buffer; @@ -105,7 +103,7 @@ pub struct MutableBuffer { /// Memory reservation for tracking memory usage #[cfg(feature = "pool")] - reservation: Mutex>>, + reservation: TrackedReservation, } impl MutableBuffer { @@ -145,7 +143,7 @@ impl MutableBuffer { len: 0, layout, #[cfg(feature = "pool")] - reservation: std::sync::Mutex::new(None), + reservation: TrackedReservation::default(), } } @@ -179,7 +177,7 @@ impl MutableBuffer { len, layout, #[cfg(feature = "pool")] - reservation: std::sync::Mutex::new(None), + reservation: TrackedReservation::default(), } } @@ -193,7 +191,8 @@ impl MutableBuffer { let len = bytes.len(); let data = bytes.ptr(); #[cfg(feature = "pool")] - let reservation = bytes.reservation.lock().unwrap().take(); + let reservation = bytes.reservation.take(); + mem::forget(bytes); Ok(Self { @@ -201,7 +200,7 @@ impl MutableBuffer { len, layout, #[cfg(feature = "pool")] - reservation: Mutex::new(reservation), + reservation, }) } @@ -392,11 +391,7 @@ impl MutableBuffer { self.data = NonNull::new(data).unwrap_or_else(|| handle_alloc_error(new_layout)); self.layout = new_layout; #[cfg(feature = "pool")] - { - if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { - reservation.resize(self.layout.size()); - } - } + self.reservation.resize(self.layout.size()); } /// Truncates this buffer to `len` bytes @@ -409,11 +404,7 @@ impl MutableBuffer { } self.len = len; #[cfg(feature = "pool")] - { - if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { - reservation.resize(self.len); - } - } + self.reservation.resize(self.len); } /// Resizes the buffer, either truncating its contents (with no change in capacity), or @@ -442,12 +433,9 @@ impl MutableBuffer { } // this truncates the buffer when new_len < self.len self.len = new_len; + #[cfg(feature = "pool")] - { - if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { - reservation.resize(self.len); - } - } + self.reservation.resize(self.len); } /// Shrinks the capacity of the buffer as much as possible. @@ -502,11 +490,7 @@ impl MutableBuffer { pub fn clear(&mut self) { self.len = 0; #[cfg(feature = "pool")] - { - if let Some(reservation) = self.reservation.lock().unwrap().as_mut() { - reservation.resize(self.len); - } - } + self.reservation.resize(self.len); } /// Returns the data stored in this buffer as a slice. @@ -537,10 +521,7 @@ impl MutableBuffer { pub(super) fn into_buffer(self) -> Buffer { let bytes = unsafe { Bytes::new(self.data, self.len, Deallocation::Standard(self.layout)) }; #[cfg(feature = "pool")] - { - let reservation = self.reservation.lock().unwrap().take(); - *bytes.reservation.lock().unwrap() = reservation; - } + bytes.reservation.replace(self.reservation.take()); std::mem::forget(self); Buffer::from(bytes) } @@ -845,7 +826,7 @@ impl MutableBuffer { /// multiple arrays. #[cfg(feature = "pool")] pub fn claim(&self, pool: &dyn MemoryPool) { - *self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity())); + self.reservation.claim(pool, self.capacity()); } } @@ -892,7 +873,7 @@ impl From> for MutableBuffer { len, layout, #[cfg(feature = "pool")] - reservation: std::sync::Mutex::new(None), + reservation: TrackedReservation::default(), } } } diff --git a/arrow-buffer/src/bytes.rs b/arrow-buffer/src/bytes.rs index a80a347fa17a..6429231694f5 100644 --- a/arrow-buffer/src/bytes.rs +++ b/arrow-buffer/src/bytes.rs @@ -25,11 +25,8 @@ use std::{fmt::Debug, fmt::Formatter}; use crate::alloc::Deallocation; use crate::buffer::dangling_ptr; - -#[cfg(feature = "pool")] -use crate::pool::{MemoryPool, MemoryReservation}; #[cfg(feature = "pool")] -use std::sync::Mutex; +use crate::{MemoryPool, TrackedReservation}; /// A continuous, fixed-size, immutable memory region that knows how to de-allocate itself. /// @@ -57,7 +54,7 @@ pub(crate) struct Bytes { /// Memory reservation for tracking memory usage #[cfg(feature = "pool")] - pub(super) reservation: Mutex>>, + pub(super) reservation: TrackedReservation, } impl Bytes { @@ -80,7 +77,7 @@ impl Bytes { len, deallocation, #[cfg(feature = "pool")] - reservation: Mutex::new(None), + reservation: TrackedReservation::default(), } } @@ -110,7 +107,7 @@ impl Bytes { /// Register this [`Bytes`] with the provided [`MemoryPool`], replacing any prior reservation. #[cfg(feature = "pool")] pub(crate) fn claim(&self, pool: &dyn MemoryPool) { - *self.reservation.lock().unwrap() = Some(pool.reserve(self.capacity())); + self.reservation.claim(pool, self.capacity()); } /// Resize the memory reservation of this buffer @@ -118,14 +115,7 @@ impl Bytes { /// This is a no-op if this buffer doesn't have a reservation. #[cfg(feature = "pool")] fn resize_reservation(&self, new_size: usize) { - let mut guard = self.reservation.lock().unwrap(); - if let Some(mut reservation) = guard.take() { - // Resize the reservation - reservation.resize(new_size); - - // Put it back - *guard = Some(reservation); - } + self.reservation.resize(new_size); } /// Try to reallocate the underlying memory region to a new size (smaller or larger). @@ -234,7 +224,7 @@ impl From for Bytes { ptr: NonNull::new(value.as_ptr() as _).unwrap(), deallocation: Deallocation::Custom(std::sync::Arc::new(value), len), #[cfg(feature = "pool")] - reservation: Mutex::new(None), + reservation: TrackedReservation::default(), } } } diff --git a/arrow-buffer/src/pool.rs b/arrow-buffer/src/pool.rs index 95bd308a35be..9a305235c226 100644 --- a/arrow-buffer/src/pool.rs +++ b/arrow-buffer/src/pool.rs @@ -23,15 +23,15 @@ //! is as follows: //! //! ```text -//! (pool tracker) (resizable) +//! (pool tracker) (resizable) //! ┌──────────────────┐ fn reserve() ┌─────────────────────────┐ -//! │ trait MemoryPool │─────────────►│ trait MemoryReservation │ +//! │ trait MemoryPool │────────────►│ trait MemoryReservation │ //! └──────────────────┘ └─────────────────────────┘ //! ``` use std::fmt::Debug; -use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex}; /// A memory reservation within a [`MemoryPool`] that is freed on drop pub trait MemoryReservation: Debug + Send + Sync { @@ -52,20 +52,20 @@ pub trait MemoryReservation: Debug + Send + Sync { /// tell if the buffer is shared or not. /// /// ```text -/// Array A Array B +/// Array A Array B /// ┌────────────┐ ┌────────────┐ /// │ slices... │ │ slices... │ /// │────────────│ │────────────│ /// │ Arc │ │ Arc │ (shared buffer) -/// └─────▲──────┘ └───────▲────┘ -/// │ │ -/// │ Bytes │ -/// │ ┌─────────────┐ │ -/// │ │ data... │ │ -/// │ │─────────────│ │ -/// └──│ Memory │──┘ (tracked with a memory pool) -/// │ Reservation │ -/// └─────────────┘ +/// └─────▲─────┘ └───────▲───┘ +/// │ │ +/// │ Bytes │ +/// │ ┌─────────────┐ │ +/// │ │ data... │ │ +/// │ │─────────────│ │ +/// └──│ Memory │──┘ (tracked with a memory pool) +/// │ Reservation │ +/// └─────────────┘ /// ``` /// /// With a memory pool, we can count the memory usage by the shared buffer @@ -147,6 +147,77 @@ impl MemoryReservation for Tracker { } } +/// This is a wrapper for the reservation so we can standardize on changing +/// and avoid race conditions in memory accounting +#[derive(Debug, Default)] +pub struct TrackedReservation { + reservation: Mutex>>, +} + +impl TrackedReservation { + /// Claim memory from a pool, replacing the current reservation (if exists). + pub fn claim(&self, pool: &dyn MemoryPool, capacity: usize) { + // get the existing reservation + let mut guard = self + .reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + // drop it before we reserve the new one + drop(guard.take()); + + // reserve the new one + *guard = Some(pool.reserve(capacity)) + } + + /// Resize the memory reservation of this buffer + /// + /// This is a no-op if this buffer doesn't have a reservation. + pub fn resize(&self, new_size: usize) { + let mut guard = self + .reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + if let Some(reservation) = guard.as_mut() { + // Resize the reservation + reservation.resize(new_size); + } + } + + /// Takes ownership of the reservation and returns it in a new `TrackedReservation` + pub fn take(&self) -> Self { + let mut guard = self + .reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + let reservation = guard.take(); + + Self { + reservation: Mutex::new(reservation), + } + } + + /// Replaces the current tracked reservation with `other`, consuming it. + pub fn replace(&self, other: Self) { + // get the owned value out, preventing double lock + let reservation = other + .reservation + .into_inner() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + let mut guard = self + .reservation + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + // drop the old reservation before installing the new one + drop(guard.take()); + *guard = reservation; + } +} + #[cfg(test)] mod tests { use super::*; @@ -186,4 +257,54 @@ mod tests { drop(reservation2); assert_eq!(pool.used(), 0); } + + /// A [`MemoryPool`] that records the peak usage observed at the instant + /// each reservation is taken, letting a single-threaded test witness the + /// transient double-count that [`TrackedReservation::claim`] must avoid. + #[derive(Debug, Default)] + struct PeakPool { + inner: TrackingMemoryPool, + peak: AtomicUsize, + } + + impl MemoryPool for PeakPool { + fn reserve(&self, size: usize) -> Box { + let reservation = self.inner.reserve(size); + self.peak.fetch_max(self.inner.used(), Ordering::Relaxed); + reservation + } + + fn available(&self) -> isize { + self.inner.available() + } + + fn used(&self) -> usize { + self.inner.used() + } + + fn capacity(&self) -> usize { + self.inner.capacity() + } + } + + #[test] + fn test_claim_reclaims_before_reserving() { + let pool = PeakPool::default(); + let reservation = TrackedReservation::default(); + + // Claim 512 bytes. + reservation.claim(&pool, 512); + assert_eq!(pool.used(), 512); + + // Re-claim the same amount. The old reservation must be released + // before the new one is taken, so usage never transiently doubles + // (see #10139). + reservation.claim(&pool, 512); + assert_eq!(pool.used(), 512); + assert_eq!( + pool.peak.load(Ordering::Relaxed), + 512, + "claim double-counted memory while reclaiming" + ); + } } From 65f724a9c78326d038e87062399a5d03cbd54e20 Mon Sep 17 00:00:00 2001 From: cetra3 Date: Tue, 21 Jul 2026 15:39:32 +0800 Subject: [PATCH 2/2] Make crate visibility only --- arrow-buffer/src/pool.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/arrow-buffer/src/pool.rs b/arrow-buffer/src/pool.rs index 9a305235c226..82abd075284c 100644 --- a/arrow-buffer/src/pool.rs +++ b/arrow-buffer/src/pool.rs @@ -150,7 +150,7 @@ impl MemoryReservation for Tracker { /// This is a wrapper for the reservation so we can standardize on changing /// and avoid race conditions in memory accounting #[derive(Debug, Default)] -pub struct TrackedReservation { +pub(crate) struct TrackedReservation { reservation: Mutex>>, }