This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new 6e728ced77 Fix pool claim race condition (#10301)
6e728ced77 is described below
commit 6e728ced77368fa8a947ac52e7efba19fb5dfde9
Author: Peter L <[email protected]>
AuthorDate: Wed Sep 2 12:29:04 2026 +0930
Fix pool claim race condition (#10301)
# Which issue does this PR close?
- Closes https://github.com/apache/arrow-rs/issues/10139
# Rationale for this change
This adjusts the internals of the memory pool reservation so that there
isn't any double accounting when replacing memory with `claim`
# What changes are included in this PR?
Adds a new struct in `pool.rs` which is used to track reservations on
buffers. This is essentially a newtype lift of the existing structure,
but with a few specialised methods to ensure that claim accounts
correctly.
# Are these changes tested?
Yes a new test has been added
# Are there any user-facing changes?
No this is just some memory pool accounting internals
---------
Co-authored-by: Jeffrey Vo <[email protected]>
---
arrow-buffer/src/buffer/mutable.rs | 48 +++++---------
arrow-buffer/src/bytes.rs | 22 ++-----
arrow-buffer/src/pool.rs | 125 ++++++++++++++++++++++++++++++++++---
3 files changed, 135 insertions(+), 60 deletions(-)
diff --git a/arrow-buffer/src/buffer/mutable.rs
b/arrow-buffer/src/buffer/mutable.rs
index c53cf3fcab..7ae7d1337f 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, lock_reservation};
-#[cfg(feature = "pool")]
-use std::sync::Mutex;
+use crate::pool::{MemoryPool, TrackedReservation};
use super::Buffer;
@@ -130,7 +128,7 @@ pub struct MutableBuffer {
/// Memory reservation for tracking memory usage
#[cfg(feature = "pool")]
- reservation: Mutex<Option<Box<dyn MemoryReservation>>>,
+ reservation: TrackedReservation,
}
impl MutableBuffer {
@@ -181,7 +179,7 @@ impl MutableBuffer {
len: 0,
layout,
#[cfg(feature = "pool")]
- reservation: std::sync::Mutex::new(None),
+ reservation: TrackedReservation::default(),
})
}
@@ -224,7 +222,7 @@ impl MutableBuffer {
len,
layout,
#[cfg(feature = "pool")]
- reservation: std::sync::Mutex::new(None),
+ reservation: TrackedReservation::default(),
})
}
@@ -238,7 +236,8 @@ impl MutableBuffer {
let len = bytes.len();
let data = bytes.ptr();
#[cfg(feature = "pool")]
- let reservation = lock_reservation(&bytes.reservation).take();
+ let reservation = bytes.reservation.take();
+
mem::forget(bytes);
Ok(Self {
@@ -246,7 +245,7 @@ impl MutableBuffer {
len,
layout,
#[cfg(feature = "pool")]
- reservation: Mutex::new(reservation),
+ reservation,
})
}
@@ -446,11 +445,7 @@ impl MutableBuffer {
};
self.layout = new_layout;
#[cfg(feature = "pool")]
- {
- if let Some(reservation) =
lock_reservation(&self.reservation).as_mut() {
- reservation.resize(self.layout.size());
- }
- }
+ self.reservation.resize(self.layout.size());
Ok(())
}
/// Truncates this buffer to `len` bytes
@@ -463,11 +458,7 @@ impl MutableBuffer {
}
self.len = len;
#[cfg(feature = "pool")]
- {
- if let Some(reservation) =
lock_reservation(&self.reservation).as_mut() {
- reservation.resize(self.len);
- }
- }
+ self.reservation.resize(self.len);
}
/// Fallible version of [`MutableBuffer::resize`].
@@ -483,11 +474,7 @@ impl MutableBuffer {
// this truncates the buffer when new_len < self.len
self.len = new_len;
#[cfg(feature = "pool")]
- {
- if let Some(reservation) =
lock_reservation(&self.reservation).as_mut() {
- reservation.resize(self.len);
- }
- }
+ self.reservation.resize(self.len);
Ok(())
}
/// Resizes the buffer, either truncating its contents (with no change in
capacity), or
@@ -572,11 +559,7 @@ impl MutableBuffer {
pub fn clear(&mut self) {
self.len = 0;
#[cfg(feature = "pool")]
- {
- if let Some(reservation) =
lock_reservation(&self.reservation).as_mut() {
- reservation.resize(self.len);
- }
- }
+ self.reservation.resize(self.len);
}
/// Returns the data stored in this buffer as a slice.
@@ -607,10 +590,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 = lock_reservation(&self.reservation).take();
- *lock_reservation(&bytes.reservation) = reservation;
- }
+ bytes.reservation.replace(self.reservation.take());
std::mem::forget(self);
Buffer::from(bytes)
}
@@ -934,7 +914,7 @@ impl MutableBuffer {
/// multiple arrays.
#[cfg(feature = "pool")]
pub fn claim(&self, pool: &dyn MemoryPool) {
- *lock_reservation(&self.reservation) =
Some(pool.reserve(self.capacity()));
+ self.reservation.claim(pool, self.capacity());
}
}
@@ -981,7 +961,7 @@ impl<T: ArrowNativeType> From<Vec<T>> 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 de9f7befe6..ef3ede7e43 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, lock_reservation};
#[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,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) {
- *lock_reservation(&self.reservation) =
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 = lock_reservation(&self.reservation);
- 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).
@@ -236,7 +226,7 @@ impl From<bytes::Bytes> for Bytes {
ptr: NonNull::new(value.as_ptr().cast_mut()).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 6acae3ffdb..4f174bf992 100644
--- a/arrow-buffer/src/pool.rs
+++ b/arrow-buffer/src/pool.rs
@@ -23,7 +23,7 @@
//! is as follows:
//!
//! ```text
-//! (pool tracker) (resizable)
+//! (pool tracker) (resizable)
//! ┌──────────────────┐ fn reserve() ┌─────────────────────────┐
//! │ trait MemoryPool │─────────────►│ trait MemoryReservation │
//! └──────────────────┘ └─────────────────────────┘
@@ -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
@@ -158,6 +158,61 @@ pub(crate) fn lock_reservation(
reservation.lock().unwrap_or_else(PoisonError::into_inner)
}
+/// 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 = lock_reservation(&self.reservation);
+
+ // 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) {
+ if let Some(reservation) =
lock_reservation(&self.reservation).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 reservation = lock_reservation(&self.reservation).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(PoisonError::into_inner);
+
+ let mut guard = lock_reservation(&self.reservation);
+
+ // drop the old reservation before installing the new one
+ drop(guard.take());
+ *guard = reservation;
+ }
+}
+
#[cfg(test)]
mod tests {
use super::*;
@@ -197,4 +252,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"
+ );
+ }
}
