Add pool_size method to MemoryPool (#218) (#230)

* Add pool_size method to MemoryPool

* Fix

* Fmt

Co-authored-by: Daniël Heres <[email protected]>

Author: 2 people

datafusion/execution/src/memory_pool/mod.rs

@@ -106,6 +106,9 @@ pub trait MemoryPool: Send + Sync + std::fmt::Debug {
 
     /// Return the total amount of memory reserved
     fn reserved(&self) -> usize;
+
+    /// Return the configured pool size (if any)
+    fn pool_size(&self) -> Option<usize>;
 }
 
 /// A memory consumer is a named allocation traced by a particular
@@ -321,7 +324,9 @@ mod tests {
 
     #[test]
     fn test_memory_pool_underflow() {
-        let pool = Arc::new(GreedyMemoryPool::new(50)) as _;
+        let pool: Arc<dyn MemoryPool> = Arc::new(GreedyMemoryPool::new(50)) as _;
+        assert_eq!(pool.pool_size(), Some(50));
+
         let mut a1 = MemoryConsumer::new("a1").register(&pool);
         assert_eq!(pool.reserved(), 0);
 

datafusion/execution/src/memory_pool/pool.rs

@@ -44,6 +44,10 @@ impl MemoryPool for UnboundedMemoryPool {
     fn reserved(&self) -> usize {
         self.used.load(Ordering::Relaxed)
     }
+
+    fn pool_size(&self) -> Option<usize> {
+        None
+    }
 }
 
 /// A [`MemoryPool`] that implements a greedy first-come first-serve limit.
@@ -96,6 +100,10 @@ impl MemoryPool for GreedyMemoryPool {
     fn reserved(&self) -> usize {
         self.used.load(Ordering::Relaxed)
     }
+
+    fn pool_size(&self) -> Option<usize> {
+        Some(self.pool_size)
+    }
 }
 
 /// A [`MemoryPool`] that prevents spillable reservations from using more than
@@ -229,6 +237,10 @@ impl MemoryPool for FairSpillPool {
         let state = self.state.lock();
         state.spillable + state.unspillable
     }
+
+    fn pool_size(&self) -> Option<usize> {
+        Some(self.pool_size)
+    }
 }
 
 fn insufficient_capacity_err(
@@ -246,7 +258,8 @@ mod tests {
 
     #[test]
     fn test_fair() {
-        let pool = Arc::new(FairSpillPool::new(100)) as _;
+        let pool: Arc<dyn MemoryPool> = Arc::new(FairSpillPool::new(100)) as _;
+        assert_eq!(pool.pool_size(), Some(100));
 
         let mut r1 = MemoryConsumer::new("unspillable").register(&pool);
         // Can grow beyond capacity of pool