Rich-T-kid commented on code in PR #10301:
URL: https://github.com/apache/arrow-rs/pull/10301#discussion_r3608845225
##########
arrow-buffer/src/bytes.rs:
##########
@@ -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);
- }
Review Comment:
nice
##########
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 │
Review Comment:
Nit: there are a couple of spots in this PR where nothing has actually
changed but a diff is still being shown. We should generally try to avoid that.
##########
arrow-buffer/src/pool.rs:
##########
@@ -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() {
Review Comment:
nice regression test
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
