From 050e4021a9a5d680b2fc7c3f6708b1050376044e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20Heres?= <danielheres@gmail.com>
Date: Thu, 30 Nov 2023 17:53:11 +0000
Subject: [PATCH] Add pool_size method to MemoryPool (#218)

* Add pool_size method to MemoryPool

* Fix

* Fmt
---
 datafusion/execution/src/memory_pool/mod.rs  |  7 ++++++-
 datafusion/execution/src/memory_pool/pool.rs | 15 ++++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs
index 7816f15bc259..aeb3d179afa0 100644
--- a/datafusion/execution/src/memory_pool/mod.rs
+++ b/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);
 
diff --git a/datafusion/execution/src/memory_pool/pool.rs b/datafusion/execution/src/memory_pool/pool.rs
index 4a491630fe20..bd9f818a7aad 100644
--- a/datafusion/execution/src/memory_pool/pool.rs
+++ b/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