Fix pool claim race condition

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<Option<Box<dyn MemoryReservation>>>,
+    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,15 +191,16 @@ 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 {
             data,
             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<T: ArrowNativeType> From<Vec<T>> for MutableBuffer {
             len,
             layout,
             #[cfg(feature = "pool")]
-            reservation: std::sync::Mutex::new(None),
+            reservation: TrackedReservation::default(),
         }
     }
 }

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<Option<Box<dyn MemoryReservation>>>,
+    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,22 +107,15 @@ 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
     ///
     /// 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<bytes::Bytes> 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(),
         }
     }
 }

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<Bytes> │    │ Arc<Bytes> │ (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(crate) struct TrackedReservation {
+    reservation: Mutex<Option<Box<dyn MemoryReservation>>>,
+}
+
+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<dyn MemoryReservation> {
+            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"
+        );
+    }
 }