tricord 0.1.0

//! Periodic sampling of a process tree's resource use.
//!
//! The sampler runs on its own thread, polling the OS at a configurable
//! interval (default 0.5 s, matching Snakemake's `BENCHMARK_INTERVAL_SHORT`).
//! At each tick it asks the platform module ([`crate::platform`]) for a
//! snapshot of every live process in the tree rooted at the spawned child,
//! then folds the snapshot into a running [`SamplerState`]. When the sampler
//! is asked to stop (via [`SamplerHandle::stop`]) it returns the aggregated
//! [`BenchmarkRecord`].

// USS, PSS, RSS, VMS are domain acronyms for distinct memory metrics that
// happen to look similar to clippy's lexical similarity check. The
// distinctions matter, so we keep the names.
#![allow(clippy::similar_names)]

use std::{
    collections::HashMap,
    sync::{
        Arc,
        atomic::{AtomicBool, Ordering},
    },
    thread::{self, JoinHandle},
    time::{Duration, Instant},
};

use crate::{platform, record::BenchmarkRecord};

/// Default sampling interval (matches `snakemake.benchmark.BENCHMARK_INTERVAL_SHORT`).
pub const DEFAULT_INTERVAL: Duration = Duration::from_millis(500);

/// One snapshot of resource use for one process at one moment in time.
///
/// All byte-valued fields are raw bytes (the [`SamplerState`] aggregator
/// converts to MiB when it produces the final record).
#[derive(Debug, Clone, Copy, Default)]
pub struct ProcessSnapshot {
    pub pid: i32,
    pub rss_bytes: u64,
    pub vms_bytes: u64,
    pub uss_bytes: Option<u64>,
    pub pss_bytes: Option<u64>,
    pub io_read_bytes: Option<u64>,
    pub io_write_bytes: Option<u64>,
    pub cpu_time_seconds: f64,
}

/// Per-PID accumulator; used to keep the latest seen value of monotonically-
/// increasing counters even after the process exits.
#[derive(Debug, Clone, Copy, Default)]
struct ProcessAccum {
    io_read_bytes: Option<u64>,
    io_write_bytes: Option<u64>,
    cpu_time_seconds: f64,
}

/// Running aggregate of all snapshots taken during a benchmark run.
///
/// `Option<u64>` peak fields use `None` to mean "this metric was never observed
/// for any process in any tick" — used downstream to render the TSV column as
/// `-` rather than `0.00`.
#[derive(Debug, Default)]
pub struct SamplerState {
    max_rss_bytes: u64,
    max_vms_bytes: u64,
    max_uss_bytes: Option<u64>,
    max_pss_bytes: Option<u64>,
    per_pid: HashMap<i32, ProcessAccum>,
    data_collected: bool,
}

impl SamplerState {
    /// Fold one tick's worth of snapshots into the running aggregate.
    pub fn absorb(&mut self, snapshots: &[ProcessSnapshot]) {
        if snapshots.is_empty() {
            return;
        }
        let mut tick_rss: u64 = 0;
        let mut tick_vms: u64 = 0;
        let mut tick_uss_total: u64 = 0;
        let mut tick_pss_total: u64 = 0;
        let mut any_uss = false;
        let mut any_pss = false;

        for snap in snapshots {
            tick_rss = tick_rss.saturating_add(snap.rss_bytes);
            tick_vms = tick_vms.saturating_add(snap.vms_bytes);
            if let Some(v) = snap.uss_bytes {
                tick_uss_total = tick_uss_total.saturating_add(v);
                any_uss = true;
            }
            if let Some(v) = snap.pss_bytes {
                tick_pss_total = tick_pss_total.saturating_add(v);
                any_pss = true;
            }

            let entry = self.per_pid.entry(snap.pid).or_default();
            if let Some(io_in) = snap.io_read_bytes {
                entry.io_read_bytes = Some(io_in.max(entry.io_read_bytes.unwrap_or(0)));
            }
            if let Some(io_out) = snap.io_write_bytes {
                entry.io_write_bytes = Some(io_out.max(entry.io_write_bytes.unwrap_or(0)));
            }
            if snap.cpu_time_seconds > entry.cpu_time_seconds {
                entry.cpu_time_seconds = snap.cpu_time_seconds;
            }
        }

        self.max_rss_bytes = self.max_rss_bytes.max(tick_rss);
        self.max_vms_bytes = self.max_vms_bytes.max(tick_vms);
        if any_uss {
            self.max_uss_bytes = Some(self.max_uss_bytes.unwrap_or(0).max(tick_uss_total));
        }
        if any_pss {
            self.max_pss_bytes = Some(self.max_pss_bytes.unwrap_or(0).max(tick_pss_total));
        }
        self.data_collected = true;
    }

    /// Materialize the running aggregate into a [`BenchmarkRecord`] given the
    /// final wall-clock running time.
    #[must_use]
    pub fn into_record(self, running_time_seconds: f64) -> BenchmarkRecord {
        let mib = |b: u64| (b as f64) / (1024.0 * 1024.0);

        if !self.data_collected {
            return BenchmarkRecord {
                running_time: running_time_seconds,
                data_collected: false,
                ..Default::default()
            };
        }

        let mut io_in_bytes: u64 = 0;
        let mut io_out_bytes: u64 = 0;
        let mut cpu_time_seconds = 0.0_f64;
        let mut any_io = false;
        for accum in self.per_pid.values() {
            if let Some(v) = accum.io_read_bytes {
                io_in_bytes = io_in_bytes.saturating_add(v);
                any_io = true;
            }
            if let Some(v) = accum.io_write_bytes {
                io_out_bytes = io_out_bytes.saturating_add(v);
                any_io = true;
            }
            cpu_time_seconds += accum.cpu_time_seconds;
        }

        let mean_load = if running_time_seconds > 0.0 {
            (cpu_time_seconds / running_time_seconds) * 100.0
        } else {
            0.0
        };

        BenchmarkRecord {
            running_time: running_time_seconds,
            max_rss: Some(mib(self.max_rss_bytes)),
            max_vms: Some(mib(self.max_vms_bytes)),
            max_uss: self.max_uss_bytes.map(mib),
            max_pss: self.max_pss_bytes.map(mib),
            io_in: if any_io { Some(mib(io_in_bytes)) } else { None },
            io_out: if any_io { Some(mib(io_out_bytes)) } else { None },
            mean_load,
            cpu_time: cpu_time_seconds,
            data_collected: true,
        }
    }
}

/// Options controlling sampler behavior.
#[derive(Debug, Clone, Copy)]
pub struct SamplerOptions {
    /// Wall-clock interval between samples.
    pub interval: Duration,
}

impl Default for SamplerOptions {
    fn default() -> Self {
        Self { interval: DEFAULT_INTERVAL }
    }
}

/// Handle to a running sampler thread.
pub struct SamplerHandle {
    stop: Arc<AtomicBool>,
    thread: Option<JoinHandle<SamplerState>>,
    started_at: Instant,
}

impl SamplerHandle {
    /// Spawn a background thread that polls the process tree rooted at
    /// `root_pid` until [`Self::stop`] is called.
    ///
    /// # Panics
    /// Panics if the OS refuses to spawn a new thread (extreme resource
    /// exhaustion). Callers are expected to terminate in that case.
    #[must_use]
    pub fn spawn(root_pid: i32, options: SamplerOptions) -> Self {
        let stop = Arc::new(AtomicBool::new(false));
        let stop_for_thread = Arc::clone(&stop);
        let started_at = Instant::now();
        let thread = thread::Builder::new()
            .name("tricord-sampler".into())
            .spawn(move || run_sampler_loop(root_pid, options.interval, &stop_for_thread))
            .expect("failed to spawn sampler thread");
        Self { stop, thread: Some(thread), started_at }
    }

    /// Signal the sampler thread to stop and wait for it to finish.
    ///
    /// Returns the aggregated [`BenchmarkRecord`] including the final wall-clock
    /// running time computed from when [`Self::spawn`] was called.
    ///
    /// # Panics
    /// Panics if called twice on the same handle, which the type system already
    /// prevents (the method takes `self` by value).
    #[must_use]
    pub fn stop(mut self) -> BenchmarkRecord {
        self.stop.store(true, Ordering::SeqCst);
        let state =
            self.thread.take().expect("sampler thread already joined").join().unwrap_or_default();
        let elapsed = self.started_at.elapsed().as_secs_f64();
        state.into_record(elapsed)
    }
}

fn run_sampler_loop(root_pid: i32, interval: Duration, stop: &AtomicBool) -> SamplerState {
    let mut state = SamplerState::default();
    let mut sampler = platform::new_sampler();
    while !stop.load(Ordering::SeqCst) {
        thread::sleep(interval);
        if stop.load(Ordering::SeqCst) {
            // Skip the sample if we were asked to stop while sleeping; the
            // child has already exited and re-reading /proc may race.
            break;
        }
        let snapshots = sampler.sample_tree(root_pid);
        state.absorb(&snapshots);
    }
    state
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn empty_state_records_no_data() {
        let state = SamplerState::default();
        let record = state.into_record(1.0);
        assert!(!record.data_collected);
        assert!((record.running_time - 1.0).abs() < 1e-9);
        assert!(record.max_rss.is_none());
    }

    #[test]
    fn single_snapshot_populates_record() {
        let mut state = SamplerState::default();
        state.absorb(&[ProcessSnapshot {
            pid: 100,
            rss_bytes: 10 * 1024 * 1024,
            vms_bytes: 100 * 1024 * 1024,
            uss_bytes: Some(8 * 1024 * 1024),
            pss_bytes: Some(9 * 1024 * 1024),
            io_read_bytes: Some(1024 * 1024),
            io_write_bytes: Some(2 * 1024 * 1024),
            cpu_time_seconds: 0.5,
        }]);
        let record = state.into_record(2.0);
        assert!(record.data_collected);
        assert_eq!(record.max_rss, Some(10.0));
        assert_eq!(record.max_vms, Some(100.0));
        assert_eq!(record.max_uss, Some(8.0));
        assert_eq!(record.max_pss, Some(9.0));
        assert_eq!(record.io_in, Some(1.0));
        assert_eq!(record.io_out, Some(2.0));
        assert!((record.cpu_time - 0.5).abs() < 1e-9);
        assert!((record.mean_load - 25.0).abs() < 1e-9); // 0.5s cpu / 2.0s wall = 25%.
    }

    #[test]
    fn memory_max_uses_summed_tree() {
        let mut state = SamplerState::default();
        // Tick 1: parent + child, total = 30 MiB RSS.
        state.absorb(&[
            ProcessSnapshot {
                pid: 1,
                rss_bytes: 10 * 1024 * 1024,
                vms_bytes: 0,
                cpu_time_seconds: 0.0,
                ..Default::default()
            },
            ProcessSnapshot {
                pid: 2,
                rss_bytes: 20 * 1024 * 1024,
                vms_bytes: 0,
                cpu_time_seconds: 0.0,
                ..Default::default()
            },
        ]);
        // Tick 2: parent only, total = 25 MiB RSS — peak of this tick is lower.
        state.absorb(&[ProcessSnapshot {
            pid: 1,
            rss_bytes: 25 * 1024 * 1024,
            vms_bytes: 0,
            cpu_time_seconds: 0.0,
            ..Default::default()
        }]);
        let record = state.into_record(1.0);
        assert_eq!(record.max_rss, Some(30.0)); // peak across snapshots
    }

    #[test]
    fn io_and_cpu_aggregate_across_pids_after_exit() {
        let mut state = SamplerState::default();
        // Tick 1: child A and child B both alive with I/O and CPU usage.
        state.absorb(&[
            ProcessSnapshot {
                pid: 10,
                rss_bytes: 1,
                vms_bytes: 1,
                io_read_bytes: Some(50 * 1024 * 1024),
                io_write_bytes: Some(10 * 1024 * 1024),
                cpu_time_seconds: 1.0,
                ..Default::default()
            },
            ProcessSnapshot {
                pid: 11,
                rss_bytes: 1,
                vms_bytes: 1,
                io_read_bytes: Some(20 * 1024 * 1024),
                io_write_bytes: Some(5 * 1024 * 1024),
                cpu_time_seconds: 0.5,
                ..Default::default()
            },
        ]);
        // Tick 2: child B has exited; child A continues.
        state.absorb(&[ProcessSnapshot {
            pid: 10,
            rss_bytes: 1,
            vms_bytes: 1,
            io_read_bytes: Some(60 * 1024 * 1024),
            io_write_bytes: Some(15 * 1024 * 1024),
            cpu_time_seconds: 1.5,
            ..Default::default()
        }]);
        let record = state.into_record(1.0);
        // io_in: child A latest (60) + child B latest (20) = 80 MiB
        assert_eq!(record.io_in, Some(80.0));
        assert_eq!(record.io_out, Some(20.0));
        // cpu_time: child A (1.5) + child B (0.5) = 2.0s
        assert!((record.cpu_time - 2.0).abs() < 1e-9);
    }

    #[test]
    fn missing_uss_and_io_remain_none_in_record() {
        let mut state = SamplerState::default();
        state.absorb(&[ProcessSnapshot {
            pid: 1,
            rss_bytes: 1024 * 1024,
            vms_bytes: 1024 * 1024,
            uss_bytes: None,
            pss_bytes: None,
            io_read_bytes: None,
            io_write_bytes: None,
            cpu_time_seconds: 0.0,
        }]);
        let record = state.into_record(1.0);
        assert!(record.max_uss.is_none());
        assert!(record.max_pss.is_none());
        assert!(record.io_in.is_none());
        assert!(record.io_out.is_none());
    }
}