alamb commented on code in PR #9015:
URL: https://github.com/apache/arrow-datafusion/pull/9015#discussion_r1508188518
##########
datafusion/execution/src/memory_pool/pool.rs:
##########
@@ -54,7 +58,44 @@ impl MemoryPool for UnboundedMemoryPool {
#[derive(Debug)]
pub struct GreedyMemoryPool {
pool_size: usize,
- used: AtomicUsize,
+ meta: Mutex<SimplePool>,
+}
+
+// Simple struct that tracks memory usage across multiple consumers.
+// Can be used in composition with other fields for providing more
sophisticated memory pool
+// policies
+#[derive(Debug, Default)]
+struct SimplePool {
+ // Maps consumer names to the amount of memory they use.
+ pool_members: HashMap<String, usize>,
+ // The total memory usage of all consumers.
+ total_size: usize,
Review Comment:
I think it would help the code's readiblity to use call this `used` to match
with the previous implementation
##########
datafusion/execution/src/memory_pool/pool.rs:
##########
@@ -310,4 +460,68 @@ mod tests {
let err = r4.try_grow(30).unwrap_err().strip_backtrace();
assert_eq!(err, "Resources exhausted: Failed to allocate additional 30
bytes for s4 with 0 bytes already allocated - maximum available is 20");
}
+
+ #[test]
+ fn test_greedy() {
+ let pool = Arc::new(GreedyMemoryPool::new(100)) as _;
+ let mut r1 = MemoryConsumer::new("r1").register(&pool);
+
+ // Can grow beyond capacity of pool
+ r1.grow(2000);
+ assert_eq!(pool.reserved(), 2000);
+
+ let mut r2 = MemoryConsumer::new("r2")
+ .with_can_spill(true)
+ .register(&pool);
+ // Can grow beyond capacity of pool
+ r2.grow(2000);
+
+ let err = r1.try_grow(1).unwrap_err().strip_backtrace();
+ assert_eq!(err, "Resources exhausted: Failed to allocate additional 1
bytes for r1 with 2000 bytes already allocated - maximum available is 0");
+
+ let err = r2.try_grow(1).unwrap_err().strip_backtrace();
+ assert_eq!(err, "Resources exhausted: Failed to allocate additional 1
bytes for r2 with 2000 bytes already allocated - maximum available is 0");
+
+ r1.shrink(1990);
+ r2.shrink(2000);
+
+ assert_eq!(pool.reserved(), 10);
+
+ r1.try_grow(10).unwrap();
+ assert_eq!(pool.reserved(), 20);
+ }
+
+ #[test]
+ fn test_greedy_pool_display() {
+ let greedy_pool = Arc::new(GreedyMemoryPool::new(1000));
+ let pool: Arc<dyn MemoryPool> = greedy_pool.clone();
+
+ let mut r1 = MemoryConsumer::new("r1").register(&pool);
+ let mut r2 = MemoryConsumer::new("r2").register(&pool);
+
+ r1.grow(100);
+ r2.grow(321);
+
+ assert_eq!(
+ format!("{}", greedy_pool),
+ "GreedyPool 2 allocations, 421 used, 579 free, 1000
capacity.\nConsumers:\n321: r2\n100: r1\n"
Review Comment:
👍
##########
datafusion/execution/src/memory_pool/pool.rs:
##########
@@ -163,25 +242,30 @@ impl MemoryPool for FairSpillPool {
}
fn unregister(&self, consumer: &MemoryConsumer) {
+ let mut state = self.state.lock();
+ let used = state.meta.unregister(consumer.name());
+
if consumer.can_spill {
- let mut state = self.state.lock();
- state.num_spill = state.num_spill.checked_sub(1).unwrap();
+ state.num_spill = state.num_spill.saturating_sub(1);
Review Comment:
I wonder if there is a reason to use `saturating_sub` rather than
`checked_sub`?
##########
datafusion/execution/src/memory_pool/pool.rs:
##########
@@ -124,20 +210,17 @@ impl MemoryPool for GreedyMemoryPool {
pub struct FairSpillPool {
/// The total memory limit
pool_size: usize,
-
state: Mutex<FairSpillPoolState>,
}
-#[derive(Debug)]
+#[derive(Debug, Default)]
struct FairSpillPoolState {
/// The number of consumers that can spill
num_spill: usize,
-
/// The total amount of memory reserved that can be spilled
spillable: usize,
- /// The total amount of memory reserved by consumers that cannot spill
- unspillable: usize,
+ meta: SimplePool,
Review Comment:
```suggestion
/// The total amount of memory reserved by consumers that cannot spill
meta: SimplePool,
```
--
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]
