ixa/

execution_stats.rs

use std::time::Duration;

use humantime::format_duration;
use log::info;
#[cfg(feature = "profiling")]
use log::{debug, error};
use serde_derive::Serialize;
#[cfg(feature = "profiling")]
use sysinfo::{Pid, ProcessRefreshKind, ProcessesToUpdate, System};
#[cfg(target_arch = "wasm32")]
use wasm_bindgen::{JsCast, JsValue};

#[cfg(feature = "profiling")]
use crate::profiling::QueryProfiler;

/// A monotonic timestamp used to measure profiling intervals.
///
/// Only `now()` and `elapsed()` are part of the portable API shared by native
/// and WASM targets.
#[cfg(not(target_arch = "wasm32"))]
pub type ProfilingInstant = std::time::Instant;

/// A monotonic timestamp used to measure profiling intervals.
///
/// Only `now()` and `elapsed()` are part of the portable API shared by native
/// and WASM targets.
#[cfg(target_arch = "wasm32")]
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
pub struct ProfilingInstant(f64);

#[cfg(target_arch = "wasm32")]
impl ProfilingInstant {
    #[must_use]
    pub fn now() -> Self {
        use web_sys::js_sys::Reflect;

        let global = web_sys::js_sys::global();
        let performance = Reflect::get(&global, &"performance".into())
            .ok()
            .and_then(|value: JsValue| value.dyn_into::<web_sys::Performance>().ok());

        Self(performance.map_or(0.0, |performance| performance.now()))
    }

    #[must_use]
    pub fn elapsed(&self) -> Duration {
        let elapsed_milliseconds = Self::now().0 - self.0;
        if !elapsed_milliseconds.is_finite() || elapsed_milliseconds <= 0.0 {
            return Duration::ZERO;
        }

        Duration::from_secs_f64(elapsed_milliseconds / 1000.0)
    }
}

/// A container struct for computed final statistics.
#[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,
}

#[cfg(feature = "profiling")]
#[cfg_attr(target_arch = "wasm32", allow(dead_code))]
/// How frequently we update the max memory used value.
const REFRESH_INTERVAL: Duration = Duration::from_secs(1);

#[cfg(feature = "profiling")]
#[cfg_attr(target_arch = "wasm32", allow(dead_code))]
pub(crate) struct ExecutionProfilingCollector {
    /// Aggregate profiling data for entity queries executed by this context.
    pub(crate) query_profiler: QueryProfiler,
    /// Simulation start time, used to compute elapsed wall time for the simulation execution
    start_time: ProfilingInstant,
    /// We keep track of the last time we refreshed so that client code doesn't have to and can
    /// just call `ExecutionProfilingCollector::refresh` in its event loop.
    last_refresh: ProfilingInstant,
    /// The accumulated CPU time of the process in CPU-milliseconds at simulation start, used
    /// to compute the CPU time of the simulation execution
    start_cpu_time: u64,
    /// The maximum amount of real memory used by the process as reported by
    /// `sysinfo::System::process::memory()`. This value is polled during execution to capture the
    /// max.
    max_memory_usage: u64,
    /// A `sysinfo::System` for polling memory use
    system: System,
    /// Current process, set to `None` on unsupported platforms, wasm32 in particular
    process_id: Option<Pid>,
}

#[cfg(feature = "profiling")]
impl ExecutionProfilingCollector {
    #[must_use]
    pub fn new() -> ExecutionProfilingCollector {
        let process_id = sysinfo::get_current_pid().ok();
        let now = ProfilingInstant::now();

        let mut new_stats = ExecutionProfilingCollector {
            query_profiler: QueryProfiler::default(),
            start_time: now,
            last_refresh: now,
            start_cpu_time: 0,
            max_memory_usage: 0,
            system: System::new(),
            process_id,
        };
        // Only refreshable on supported platforms.
        if let Some(process_id) = process_id {
            debug!("Process ID: {}", process_id);
            let process_refresh_kind = ProcessRefreshKind::nothing().with_cpu().with_memory();
            new_stats.update_system_info(process_refresh_kind);

            let process = new_stats.system.process(process_id).unwrap();

            new_stats.max_memory_usage = process.memory();
            new_stats.start_cpu_time = process.accumulated_cpu_time();
        }

        new_stats
    }

    /// If at least `REFRESH_INTERVAL` (1 second) has passed since the previous
    /// refresh, memory usage is polled and updated. Call this method as frequently
    /// as you like, as it takes care of limiting polling frequency itself.
    #[inline]
    pub fn refresh(&mut self) {
        #[cfg(not(target_arch = "wasm32"))]
        if self.last_refresh.elapsed() >= REFRESH_INTERVAL {
            self.poll_memory();
            self.last_refresh = ProfilingInstant::now();
        }
    }

    #[cfg_attr(target_arch = "wasm32", allow(dead_code))]
    /// Updates maximum memory usage. This method should be called about once per second,
    /// as it is a relatively expensive system call.
    fn poll_memory(&mut self) {
        if let Some(pid) = self.process_id {
            // Only refreshes memory statistics
            self.update_system_info(ProcessRefreshKind::nothing().with_memory());
            let process = self.system.process(pid).unwrap();
            self.max_memory_usage = self.max_memory_usage.max(process.memory());
        }
    }

    /// Gives accumulated CPU time of the process in CPU-milliseconds since simulation start.
    #[allow(unused)]
    #[must_use]
    pub fn cpu_time(&mut self) -> u64 {
        if let Some(process_id) = self.process_id {
            // Only refresh cpu statistics
            self.update_system_info(ProcessRefreshKind::nothing().with_cpu());

            let process = self.system.process(process_id).unwrap();
            process.accumulated_cpu_time() - self.start_cpu_time
        } else {
            0
        }
    }

    /// Refreshes the internal `sysinfo::System` object for this process using the given
    /// [`ProcessRefreshKind`](sysinfo::ProcessRefreshKind).
    #[inline]
    fn update_system_info(&mut self, process_refresh_kind: ProcessRefreshKind) {
        if let Some(pid) = self.process_id {
            if self.system.refresh_processes_specifics(
                ProcessesToUpdate::Some(&[pid]),
                true,
                process_refresh_kind,
            ) < 1
            {
                error!("could not refresh process statistics");
            }
        }
    }

    /// Computes the final summary statistics
    #[must_use]
    pub fn compute_final_statistics(&mut self) -> ExecutionStatistics {
        let mut cpu_time_millis = 0;

        if let Some(pid) = self.process_id {
            // Update both memory and cpu statistics
            self.update_system_info(ProcessRefreshKind::nothing().with_cpu().with_memory());
            let process = self.system.process(pid).unwrap();

            self.max_memory_usage = self.max_memory_usage.max(process.memory());
            cpu_time_millis = process.accumulated_cpu_time() - self.start_cpu_time;
        }

        // Convert to `Duration`s in preparation for formatting
        let cpu_time = Duration::from_millis(cpu_time_millis);
        let wall_time = self.start_time.elapsed();

        ExecutionStatistics {
            max_memory_usage: self.max_memory_usage,
            max_plans_in_flight: 0,
            max_plan_queue_memory_in_use: 0,
            cpu_time,
            wall_time,
        }
    }
}

#[cfg(not(feature = "profiling"))]
pub(crate) struct ExecutionProfilingCollector {
    start_time: ProfilingInstant,
}

#[cfg(not(feature = "profiling"))]
impl ExecutionProfilingCollector {
    #[must_use]
    pub fn new() -> ExecutionProfilingCollector {
        ExecutionProfilingCollector {
            start_time: ProfilingInstant::now(),
        }
    }

    #[inline]
    pub fn refresh(&mut self) {}

    #[must_use]
    pub fn compute_final_statistics(&mut self) -> ExecutionStatistics {
        let wall_time = self.start_time.elapsed();

        ExecutionStatistics {
            max_memory_usage: 0,
            max_plans_in_flight: 0,
            max_plan_queue_memory_in_use: 0,
            cpu_time: Duration::ZERO,
            wall_time,
        }
    }
}

/// Prints execution statistics to the console.
///
/// Use `ExecutionProfilingCollector::compute_final_statistics()` to construct [`ExecutionStatistics`].
pub fn print_execution_statistics(summary: &ExecutionStatistics) {
    println!("━━━━ Execution Summary ━━━━");
    #[cfg(feature = "profiling")]
    {
        if cfg!(target_family = "wasm") {
            println!("Memory and CPU statistics are not available on your platform.");
        } else {
            println!(
                "{:<25}{}",
                "Max memory usage:",
                bytesize::ByteSize::b(summary.max_memory_usage)
            );
            println!(
                "{:<25}{}",
                "Max plans in flight:", summary.max_plans_in_flight
            );
            println!(
                "{:<25}{}",
                "Max plan queue memory:",
                bytesize::ByteSize::b(summary.max_plan_queue_memory_in_use)
            );
            println!("{:<25}{}", "CPU time:", format_duration(summary.cpu_time));
        }
    }

    println!("{:<25}{}", "Wall time:", format_duration(summary.wall_time));
}

/// Logs execution statistics with the logging system.
///
/// Use `ExecutionProfilingCollector::compute_final_statistics()` to construct [`ExecutionStatistics`].
pub fn log_execution_statistics(stats: &ExecutionStatistics) {
    info!("Execution complete.");
    #[cfg(feature = "profiling")]
    {
        if cfg!(target_family = "wasm") {
            info!("Memory and CPU statistics are not available on your platform.");
        } else {
            info!(
                "Max memory usage: {}",
                bytesize::ByteSize::b(stats.max_memory_usage)
            );
            info!("Max plans in flight: {}", stats.max_plans_in_flight);
            info!(
                "Max plan queue memory: {}",
                bytesize::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));
}

#[cfg(all(test, feature = "profiling"))]
mod tests {
    use std::thread;
    use std::time::Duration;

    use super::*;

    #[test]
    fn test_collector_initialization() {
        let collector = ExecutionProfilingCollector::new();

        // Ensure that initial max memory usage is non-zero
        assert!(collector.max_memory_usage > 0);
    }

    #[test]
    fn test_refresh_respects_interval() {
        let mut collector = ExecutionProfilingCollector::new();
        let before = collector.max_memory_usage;

        // Call refresh immediately — it should not poll
        collector.refresh();
        let after = collector.max_memory_usage;
        assert_eq!(before, after);

        // Sleep enough time to trigger refresh
        thread::sleep(Duration::from_secs(2));
        collector.refresh();
        // Now memory usage should be refreshed — allow it to stay same or increase
        assert!(collector.max_memory_usage >= before);
    }

    #[test]
    fn test_compute_final_statistics_structure() {
        let mut collector = ExecutionProfilingCollector::new();

        thread::sleep(Duration::from_millis(100));
        let stats = collector.compute_final_statistics();

        // 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);
    }

    #[test]
    fn test_cpu_time_increases_over_time() {
        let mut collector = ExecutionProfilingCollector::new();

        // Burn ~30ms CPU time. Likely will be < 30ms, as this thread will not have 100% of CPU
        // during 30ms wall time.
        let start = ProfilingInstant::now();
        while start.elapsed().as_millis() < 30u128 {
            std::hint::black_box(0); // Prevent optimization
        }

        let cpu_time_1 = collector.cpu_time();

        // Burn ~50ms CPU time
        let start = ProfilingInstant::now();
        while start.elapsed().as_millis() < 50u128 {
            std::hint::black_box(0); // Prevent optimization
        }

        let cpu_time_2 = collector.cpu_time();
        assert!(cpu_time_2 > cpu_time_1);
    }
}