wiedld commented on code in PR #11665:
URL: https://github.com/apache/datafusion/pull/11665#discussion_r1697608547
##########
datafusion/execution/src/memory_pool/pool.rs:
##########
@@ -240,6 +249,149 @@ fn insufficient_capacity_err(
resources_datafusion_err!("Failed to allocate additional {} bytes for {}
with {} bytes already allocated - maximum available is {}", additional,
reservation.registration.consumer.name, reservation.size, available)
}
+/// A [`MemoryPool`] that tracks the consumers that have
+/// reserved memory within the inner memory pool.
+///
+/// Tracking is per hashed [`MemoryConsumer`], not per [`MemoryReservation`].
+/// The same consumer can have multiple reservations.
+#[derive(Debug)]
+pub struct TrackConsumersPool<I> {
+ inner: I,
+ top: NonZeroUsize,
+ tracked_consumers: Mutex<HashMap<MemoryConsumer, AtomicU64>>,
+}
+
+impl<I: MemoryPool> TrackConsumersPool<I> {
+ /// Creates a new [`TrackConsumersPool`].
+ ///
+ /// The `top` determines how many Top K [`MemoryConsumer`]s to include
+ /// in the reported [`DataFusionError::ResourcesExhausted`].
+ pub fn new(inner: I, top: NonZeroUsize) -> Self {
+ Self {
+ inner,
+ top,
+ tracked_consumers: Default::default(),
+ }
+ }
+
+ /// Determine if there are multiple [`MemoryConsumer`]s registered
+ /// which have the same name.
+ ///
+ /// This is very tied to the implementation of the memory consumer.
+ fn has_multiple_consumers(&self, name: &String) -> bool {
+ let consumer = MemoryConsumer::new(name);
+ let consumer_with_spill = consumer.clone().with_can_spill(true);
+ let guard = self.tracked_consumers.lock();
+ guard.contains_key(&consumer) &&
guard.contains_key(&consumer_with_spill)
+ }
+
+ /// The top consumers in a report string.
+ fn report_top(&self) -> String {
Review Comment:
The use of TrackConsumersPool for error reporting (when passed as `Arc<dyn
MemoryPool>`) is constrained by the trait definition. However, we could use the
downcasted struct itself for runtime metrics [as shown in this added
commit](https://github.com/apache/datafusion/pull/11665/commits/c8c01961dd54df1703d7efdf216a8f63db8656b4).
Is this what you were thinking?
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
