chore: log memory reservations

chore: log reservation pool at regular intervals

Author: wiedld

datafusion-cli/Cargo.lock

@@ -48,4 +48,5 @@ object_store = { workspace = true }
 parking_lot = { workspace = true }
 rand = { workspace = true }
 tempfile = { workspace = true }
+tokio = { workspace = true }
 url = { workspace = true }

datafusion/execution/Cargo.toml

@@ -16,13 +16,17 @@
 // under the License.
 
 use crate::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation};
+use chrono::Utc;
 use datafusion_common::{resources_datafusion_err, DataFusionError, Result};
 use hashbrown::HashMap;
 use log::debug;
 use parking_lot::Mutex;
 use std::{
     num::NonZeroUsize,
-    sync::atomic::{AtomicU64, AtomicUsize, Ordering},
+    sync::{
+        atomic::{AtomicU64, AtomicUsize, Ordering},
+        Arc,
+    },
 };
 
 /// A [`MemoryPool`] that enforces no limit
@@ -261,7 +265,7 @@ fn insufficient_capacity_err(
 pub struct TrackConsumersPool<I> {
     inner: I,
     top: NonZeroUsize,
-    tracked_consumers: Mutex<HashMap<MemoryConsumer, AtomicU64>>,
+    tracked_consumers: Arc<Mutex<HashMap<MemoryConsumer, AtomicU64>>>,
 }
 
 impl<I: MemoryPool> TrackConsumersPool<I> {
@@ -270,28 +274,53 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
     /// The `top` determines how many Top K [`MemoryConsumer`]s to include
     /// in the reported [`DataFusionError::ResourcesExhausted`].
     pub fn new(inner: I, top: NonZeroUsize) -> Self {
+        let tracked_consumers = Default::default();
+
+        let _captured: Arc<
+            parking_lot::lock_api::Mutex<
+                parking_lot::RawMutex,
+                HashMap<MemoryConsumer, AtomicU64>,
+            >,
+        > = Arc::clone(&tracked_consumers);
+        #[allow(clippy::disallowed_methods)]
+        tokio::spawn(async move {
+            loop {
+                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
+
+                println!(
+                    "REPORT: {}, {}",
+                    Utc::now(),
+                    Self::_report_top(5, Arc::clone(&_captured))
+                );
+            }
+        });
+
         Self {
             inner,
             top,
-            tracked_consumers: Default::default(),
+            tracked_consumers,
         }
     }
 
     /// 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 {
+    fn has_multiple_consumers(
+        name: &String,
+        tracked_consumers: &Arc<Mutex<HashMap<MemoryConsumer, AtomicU64>>>,
+    ) -> bool {
         let consumer = MemoryConsumer::new(name);
         let consumer_with_spill = consumer.clone().with_can_spill(true);
-        let guard = self.tracked_consumers.lock();
+        let guard = tracked_consumers.lock();
         guard.contains_key(&consumer) && guard.contains_key(&consumer_with_spill)
     }
 
-    /// The top consumers in a report string.
-    pub fn report_top(&self, top: usize) -> String {
-        let mut consumers = self
-            .tracked_consumers
+    fn _report_top(
+        top: usize,
+        tracked_consumers: Arc<Mutex<HashMap<MemoryConsumer, AtomicU64>>>,
+    ) -> String {
+        let mut consumers = tracked_consumers
             .lock()
             .iter()
             .map(|(consumer, reserved)| {
@@ -301,12 +330,13 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
                 )
             })
             .collect::<Vec<_>>();
-        consumers.sort_by(|a, b| b.1.cmp(&a.1)); // inverse ordering
+        // consumers.sort_by(|a, b| b.1.cmp(&a.1)); // inverse ordering
+        consumers.sort_by(|a, b| b.0.cmp(&a.0)); // sort by name, make stdout easily parsible
 
         consumers[0..std::cmp::min(top, consumers.len())]
             .iter()
             .map(|((name, can_spill), size)| {
-                if self.has_multiple_consumers(name) {
+                if Self::has_multiple_consumers(name, &tracked_consumers) {
                     format!("{name}(can_spill={}) consumed {:?} bytes", can_spill, size)
                 } else {
                     format!("{name} consumed {:?} bytes", size)
@@ -315,6 +345,11 @@ impl<I: MemoryPool> TrackConsumersPool<I> {
             .collect::<Vec<_>>()
             .join(", ")
     }
+
+    /// The top consumers in a report string.
+    pub fn report_top(&self, top: usize) -> String {
+        Self::_report_top(top, Arc::clone(&self.tracked_consumers))
+    }
 }
 
 impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
@@ -380,6 +415,7 @@ impl<I: MemoryPool> MemoryPool for TrackConsumersPool<I> {
             .and_modify(|bytes| {
                 bytes.fetch_add(additional as u64, Ordering::AcqRel);
             });
+
         Ok(())
     }