feat: Add execution statistics for max plans in flight and plan queue memory usage (#815)

Author: RobertJacobsonCDC

src/context.rs

@@ -587,7 +587,10 @@ impl Context {
     }
 
     pub fn get_execution_statistics(&mut self) -> ExecutionStatistics {
-        self.execution_profiler.compute_final_statistics()
+        let mut stats = self.execution_profiler.compute_final_statistics();
+        stats.max_plans_in_flight = self.plan_queue.max_plans_in_flight;
+        stats.max_plan_queue_memory_in_use = self.plan_queue.max_memory_in_use;
+        stats
     }
 }
 

src/execution_stats.rs

@@ -31,6 +31,8 @@ pub fn get_high_res_time() -> f64 {
 #[derive(Serialize)]
 pub struct ExecutionStatistics {
     pub max_memory_usage: u64,
+    pub max_plans_in_flight: u64,
+    pub max_plan_queue_memory_in_use: u64,
     pub cpu_time: Duration,
     pub wall_time: Duration,
 }
@@ -171,6 +173,8 @@ impl ExecutionProfilingCollector {
 
         ExecutionStatistics {
             max_memory_usage: self.max_memory_usage,
+            max_plans_in_flight: 0,
+            max_plan_queue_memory_in_use: 0,
             cpu_time,
             wall_time,
         }
@@ -190,6 +194,15 @@ pub fn print_execution_statistics(summary: &ExecutionStatistics) {
             "Max memory usage:",
             ByteSize::b(summary.max_memory_usage)
         );
+        println!(
+            "{:<25}{}",
+            "Max plans in flight:", summary.max_plans_in_flight
+        );
+        println!(
+            "{:<25}{}",
+            "Max plan queue memory:",
+            ByteSize::b(summary.max_plan_queue_memory_in_use)
+        );
         println!("{:<25}{}", "CPU time:", format_duration(summary.cpu_time));
     }
 
@@ -205,6 +218,11 @@ pub fn log_execution_statistics(stats: &ExecutionStatistics) {
         info!("Memory and CPU statistics are not available on your platform.");
     } else {
         info!("Max memory usage: {}", ByteSize::b(stats.max_memory_usage));
+        info!("Max plans in flight: {}", stats.max_plans_in_flight);
+        info!(
+            "Max plan queue memory: {}",
+            ByteSize::b(stats.max_plan_queue_memory_in_use)
+        );
         info!("CPU time: {}", format_duration(stats.cpu_time));
     }
     info!("Wall time: {}", format_duration(stats.wall_time));
@@ -251,6 +269,8 @@ mod tests {
 
         // Fields should be non-zero
         assert!(stats.max_memory_usage > 0);
+        assert_eq!(stats.max_plans_in_flight, 0);
+        assert_eq!(stats.max_plan_queue_memory_in_use, 0);
         assert!(stats.wall_time > Duration::ZERO);
     }
 

src/plan.rs

@@ -31,7 +31,14 @@ use crate::{trace, HashMap, HashMapExt};
 pub struct Queue<T, P: Eq + PartialEq + Ord> {
     queue: BinaryHeap<PlanSchedule<P>>,
     data_map: HashMap<u64, T>,
+    /// The number of plans that have been added; equivalently, the next plan ID that
+    /// will be issued.
     plan_counter: u64,
+    /// Tracks the high water mark of plans in flight (scheduled but not yet executed).
+    /// This is the max of `self.queue.len()`, not of `self.data_map.len()`.
+    pub(crate) max_plans_in_flight: u64,
+    /// Tracks the high water mark of memory allocated (capacity, not use) by this structure
+    pub(crate) max_memory_in_use: u64,
 }
 
 impl<T, P: Eq + PartialEq + Ord> Queue<T, P> {
@@ -42,6 +49,8 @@ impl<T, P: Eq + PartialEq + Ord> Queue<T, P> {
             queue: BinaryHeap::new(),
             data_map: HashMap::new(),
             plan_counter: 0,
+            max_plans_in_flight: 0,
+            max_memory_in_use: 0,
         }
     }
 
@@ -60,6 +69,11 @@ impl<T, P: Eq + PartialEq + Ord> Queue<T, P> {
         });
         self.data_map.insert(plan_id, data);
         self.plan_counter += 1;
+        self.max_plans_in_flight = self.max_plans_in_flight.max(self.queue.len() as u64);
+        self.max_memory_in_use = self
+            .max_memory_in_use
+            .max(self.estimated_memory_in_use() as u64);
+
         PlanId(plan_id)
     }
 
@@ -148,6 +162,14 @@ impl<T, P: Eq + PartialEq + Ord> Queue<T, P> {
     pub(crate) fn remaining_plan_count(&self) -> usize {
         self.queue.len()
     }
+
+    fn estimated_memory_in_use(&self) -> usize {
+        let queue_bytes = self.queue.capacity() * size_of::<PlanSchedule<P>>();
+
+        let map_entry_bytes = self.data_map.capacity() * size_of::<(u64, T)>();
+
+        queue_bytes + map_entry_bytes
+    }
 }
 
 impl<T, P: Eq + PartialEq + Ord> Default for Queue<T, P> {

src/profiling/file.rs

@@ -32,6 +32,8 @@ impl Serialize for SerializableDuration {
 #[derive(Serialize)]
 struct SerializableExecutionStatistics {
     max_memory_usage: u64,
+    max_plans_in_flight: u64,
+    max_plan_queue_memory_in_use: u64,
     cpu_time: SerializableDuration,
     wall_time: SerializableDuration,
 }
@@ -41,6 +43,8 @@ impl From<ExecutionStatistics> for SerializableExecutionStatistics {
     fn from(value: ExecutionStatistics) -> Self {
         SerializableExecutionStatistics {
             max_memory_usage: value.max_memory_usage,
+            max_plans_in_flight: value.max_plans_in_flight,
+            max_plan_queue_memory_in_use: value.max_plan_queue_memory_in_use,
             cpu_time: SerializableDuration(value.cpu_time),
             wall_time: SerializableDuration(value.wall_time),
         }
@@ -189,6 +193,8 @@ mod tests {
 
         let exec_stats = ExecutionStatistics {
             max_memory_usage: 1024 * 1024,
+            max_plans_in_flight: 12,
+            max_plan_queue_memory_in_use: 4096,
             cpu_time: Duration::from_secs(1),
             wall_time: Duration::from_secs(2),
         };
@@ -211,6 +217,11 @@ mod tests {
             json["execution_statistics"]["max_memory_usage"],
             1024 * 1024
         );
+        assert_eq!(json["execution_statistics"]["max_plans_in_flight"], 12);
+        assert_eq!(
+            json["execution_statistics"]["max_plan_queue_memory_in_use"],
+            4096
+        );
 
         let counts = json["named_counts"].as_array().unwrap();
         assert!(!counts.is_empty());
@@ -238,6 +249,8 @@ mod tests {
 
         let exec_stats = ExecutionStatistics {
             max_memory_usage: 2048,
+            max_plans_in_flight: 3,
+            max_plan_queue_memory_in_use: 8192,
             cpu_time: Duration::from_secs_f64(1.5),
             wall_time: Duration::from_secs_f64(2.5),
         };