wiedld commented on code in PR #11665:
URL: https://github.com/apache/datafusion/pull/11665#discussion_r1697466026
##########
datafusion/execution/src/memory_pool/pool.rs:
##########
@@ -311,4 +463,190 @@ 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_tracked_consumers_pool() {
+ let pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
+ GreedyMemoryPool::new(100),
+ NonZeroUsize::new(3).unwrap(),
+ ));
+
+ // Test: use all the different interfaces to change reservation size
+
+ // set r1=50, using grow and shrink
+ let mut r1 = MemoryConsumer::new("r1").register(&pool);
+ r1.grow(70);
+ r1.shrink(20);
+
+ // set r2=15 using try_grow
+ let mut r2 = MemoryConsumer::new("r2").register(&pool);
+ r2.try_grow(15)
+ .expect("should succeed in memory allotment for r2");
+
+ // set r3=20 using try_resize
+ let mut r3 = MemoryConsumer::new("r3").register(&pool);
+ r3.try_resize(25)
+ .expect("should succeed in memory allotment for r3");
+ r3.try_resize(20)
+ .expect("should succeed in memory allotment for r3");
+
+ // set r4=10
+ // this should not be reported in top 3
+ let mut r4 = MemoryConsumer::new("r4").register(&pool);
+ r4.grow(10);
+
+ // Test: reports if new reservation causes error
+ // using the previously set sizes for other consumers
+ let mut r5 = MemoryConsumer::new("r5").register(&pool);
+ let expected = "Resources exhausted with top memory consumers (across
reservations) are: r1 consumed 50 bytes, r3 consumed 20 bytes, r2 consumed 15
bytes. Error: Failed to allocate additional 150 bytes for r5 with 0 bytes
already allocated - maximum available is 5";
+ let res = r5.try_grow(150);
+ assert!(
+ matches!(
+ &res,
+ Err(DataFusionError::ResourcesExhausted(ref e)) if
e.to_string().contains(expected)
+ ),
+ "should provide list of top memory consumers, instead found {:?}",
+ res
+ );
+ }
+
+ #[test]
+ fn test_tracked_consumers_pool_register() {
+ let pool: Arc<dyn MemoryPool> = Arc::new(TrackConsumersPool::new(
+ GreedyMemoryPool::new(100),
+ NonZeroUsize::new(3).unwrap(),
+ ));
+
+ let same_name = "foo";
+
+ // Test: see error message when no consumers recorded yet
+ let mut r0 = MemoryConsumer::new(same_name).register(&pool);
+ let expected = "Resources exhausted with top memory consumers (across
reservations) are: foo consumed 0 bytes. Error: Failed to allocate additional
150 bytes for foo with 0 bytes already allocated - maximum available is 100";
+ let res = r0.try_grow(150);
+ assert!(
+ matches!(
+ &res,
+ Err(DataFusionError::ResourcesExhausted(ref e)) if
e.to_string().contains(expected)
+ ),
+ "should provide proper error when no reservations have been made
yet, instead found {:?}", res
+ );
+
+ // API: multiple registrations using the same hashed consumer,
+ // will be recognized as the same in the TrackConsumersPool.
+
+ // Test: will be the same per Top Consumers reported.
+ r0.grow(10); // make r0=10, pool available=90
+ let new_consumer_same_name = MemoryConsumer::new(same_name);
+ let mut r1 = new_consumer_same_name.clone().register(&pool);
+ // TODO: the insufficient_capacity_err() message is per reservation,
not per consumer.
+ // a followup PR will clarify this message "0 bytes already allocated
for this reservation"
+ let expected = "Resources exhausted with top memory consumers (across
reservations) are: foo consumed 10 bytes. Error: Failed to allocate additional
150 bytes for foo with 0 bytes already allocated - maximum available is 90";
+ let res = r1.try_grow(150);
+ assert!(
+ matches!(
+ &res,
+ Err(DataFusionError::ResourcesExhausted(ref e)) if
e.to_string().contains(expected)
+ ),
+ "should provide proper error with same hashed consumer (a single
foo=10 bytes, available=90), instead found {:?}", res
+ );
+
+ // Test: will accumulate size changes per consumer, not per reservation
+ r1.grow(20);
+ let expected = "Resources exhausted with top memory consumers (across
reservations) are: foo consumed 30 bytes. Error: Failed to allocate additional
150 bytes for foo with 20 bytes already allocated - maximum available is 70";
Review Comment:
I agree this is confusing. Because the error message `for foo with 20 bytes
already allocated` is [the error from the
MemoryReservation](https://github.com/apache/datafusion/blob/c3ce60f211c64e083897703adc076807c20e3617/datafusion/execution/src/memory_pool/pool.rs#L238-L244),
which has 20 bytes.
Whereas the MemoryConsumer happens to have 2 reservations -- with a total of
30 bytes allotted (across reservations).
I'm hoping the followup PR (with the message change for the
reservation-specific error) will make this more clear. Does the below read
better?
`Resources exhausted with top memory consumers (across reservations) are:
foo consumed 30 bytes. Error: Failed to allocate additional 150 bytes for foo
with 20 bytes already allocated for this reservation - 5 bytes remain available
for the total pool
`
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
