twitcher 0.7.0

//! Per-alignment execution statistics (runtime, memory, provenance) and their CSV output.
//!
//! The statistics travel with the alignment result, so they survive a round trip through the
//! in-memory cache and the alignment database: a cached hit reports the runtime and memory of the
//! computation that originally produced it, and the `source` column says where the row came from.

use std::{io::Write, sync::Mutex, time::Duration};

use generic_a_star::cost::AStarCost as _;
use lib_tsalign::a_star_aligner::template_switch_distance::AlignmentType;
use serde::{Deserialize, Serialize};
use tracing::error;

use crate::common::{
    aligner::result::{
        AlignmentFailure, SoftFailureReason, TwitcherAlignment, TwitcherAlignmentStatistics,
    },
    coords::GenomeRegion,
};

/// Where a result handed to a consumer came from.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AlignmentSource {
    /// Computed by a freshly spawned aligner process.
    Computed,
    /// Served from the in-memory cache (either finished, or awaited while in progress).
    MemCache,
    /// Loaded from the alignment database.
    Db,
}

impl AlignmentSource {
    const fn as_str(self) -> &'static str {
        match self {
            Self::Computed => "computed",
            Self::MemCache => "mem_cache",
            Self::Db => "db",
        }
    }
}

/// Peak memory of an aligner process, in bytes.
///
/// Both values are `None` whenever they could not be observed, which is the case when the process
/// died before it could report them (a hard OOM abort, for example).
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct MemoryUsage {
    /// Peak resident set size (`VmHWM`).
    pub peak_rss: Option<u64>,
    /// Peak virtual memory size (`VmPeak`). This is what the `RLIMIT_AS` of the aligner bounds, so
    /// it is the value to compare against the memory allowance.
    pub peak_vsz: Option<u64>,
}

impl MemoryUsage {
    /// Peak memory of the calling process.
    pub fn of_current_process() -> Self {
        Self::of_status_file("/proc/self/status")
    }

    /// Peak memory of a still-running process. Yields the default once the process is gone.
    pub fn of_process(pid: u32) -> Self {
        Self::of_status_file(&format!("/proc/{pid}/status"))
    }

    fn of_status_file(path: &str) -> Self {
        std::fs::read_to_string(path)
            .map(|status| parse_proc_status(&status))
            .unwrap_or_default()
    }
}

/// Extract the peak memory fields from the contents of a `/proc/<pid>/status` file.
fn parse_proc_status(status: &str) -> MemoryUsage {
    let mut memory = MemoryUsage::default();
    for line in status.lines() {
        let Some((key, value)) = line.split_once(':') else {
            continue;
        };
        let field = match key {
            "VmPeak" => &mut memory.peak_vsz,
            "VmHWM" => &mut memory.peak_rss,
            _ => continue,
        };
        *field = parse_kb(value);
    }
    memory
}

/// Parse a ` 12345 kB` value into bytes.
fn parse_kb(value: &str) -> Option<u64> {
    let value = value.trim();
    let number = value.strip_suffix("kB").unwrap_or(value).trim();
    number.parse::<u64>().ok().map(|kb| kb.saturating_mul(1024))
}

/// What it took to run one alignment, as observed by the orchestrator.
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct ExecStats {
    /// Wall clock time of the whole alignment, including spawning the aligner process.
    pub wall: Duration,
    pub memory: MemoryUsage,
    /// The memory allowance the aligner process was given, in bytes.
    pub memory_limit: usize,
}

/// The parts of a stats row that the alignment itself does not know about.
pub struct AlignmentStatsContext<'a> {
    pub cluster_region: &'a GenomeRegion,
    pub source: AlignmentSource,
    pub reference_length: usize,
    pub query_length: usize,
}

/// Appends one row per alignment request to a CSV file.
pub struct AlignmentStatsWriter {
    inner: Mutex<csv::Writer<Box<dyn Write + Send>>>,
    aligner: &'static str,
}

impl AlignmentStatsWriter {
    pub fn new(write: Box<dyn Write + Send>, aligner: &'static str) -> Self {
        Self {
            inner: Mutex::new(csv::Writer::from_writer(write)),
            aligner,
        }
    }

    pub fn write(&self, alignment: &TwitcherAlignment, context: &AlignmentStatsContext<'_>) {
        let record = AlignmentStatsRecord::new(alignment, context, self.aligner);
        let mut writer = self.inner.lock().unwrap();
        let _ = writer
            .serialize(record)
            .inspect_err(|e| error!("Can't write alignment statistics: {e}"));
    }
}

#[derive(Serialize)]
struct AlignmentStatsRecord<'a> {
    cluster_region: String,
    source: &'static str,
    aligner: &'static str,
    outcome: &'static str,
    cost: Option<u64>,
    ts_num: Option<usize>,
    wall_ms: f64,
    align_ms: Option<f64>,
    peak_rss: Option<u64>,
    peak_vsz: Option<u64>,
    memory_limit: usize,
    opened_nodes: Option<f64>,
    closed_nodes: Option<f64>,
    suboptimal_nodes: Option<f64>,
    suboptimal_ratio: Option<f64>,
    cost_per_base: Option<f64>,
    reference_length: usize,
    query_length: usize,
    error: Option<&'a str>,
}

impl<'a> AlignmentStatsRecord<'a> {
    fn new(
        alignment: &'a TwitcherAlignment,
        context: &AlignmentStatsContext<'_>,
        aligner: &'static str,
    ) -> Self {
        let mut record = Self {
            cluster_region: context.cluster_region.to_string(),
            source: context.source.as_str(),
            aligner,
            outcome: outcome(alignment),
            cost: None,
            ts_num: None,
            wall_ms: duration_ms(alignment.exec.wall),
            align_ms: None,
            peak_rss: alignment.exec.memory.peak_rss,
            peak_vsz: alignment.exec.memory.peak_vsz,
            memory_limit: alignment.exec.memory_limit,
            opened_nodes: None,
            closed_nodes: None,
            suboptimal_nodes: None,
            suboptimal_ratio: None,
            cost_per_base: None,
            reference_length: context.reference_length,
            query_length: context.query_length,
            error: None,
        };

        match &alignment.outcome {
            Ok(success) => {
                record.cost = Some(success.alignment.cost.as_primitive());
                record.ts_num = Some(
                    success
                        .alignment
                        .alignment
                        .iter_compact()
                        .filter(|(_, ty)| {
                            matches!(ty, AlignmentType::TemplateSwitchEntrance { .. })
                        })
                        .count(),
                );
                record.align_ms = Some(duration_ms(success.stats.duration()));
                if let TwitcherAlignmentStatistics::TSAlign(stats) = &success.stats {
                    record.opened_nodes = Some(stats.opened_nodes.raw());
                    record.closed_nodes = Some(stats.closed_nodes.raw());
                    record.suboptimal_nodes = Some(stats.suboptimal_opened_nodes.raw());
                    record.suboptimal_ratio = Some(stats.suboptimal_opened_nodes_ratio.raw());
                    record.cost_per_base = Some(stats.cost_per_base.raw());
                }
            }
            Err(
                AlignmentFailure::SoftFailure {
                    reason: SoftFailureReason::Other(error),
                }
                | AlignmentFailure::Error { error },
            ) => record.error = Some(error),
            Err(_) => {}
        }

        record
    }
}

fn outcome(alignment: &TwitcherAlignment) -> &'static str {
    match &alignment.outcome {
        Ok(success) if success.has_ts() => "ts",
        Ok(_) => "no_ts",
        Err(AlignmentFailure::SoftFailure {
            reason: SoftFailureReason::OutOfMemory,
        }) => "oom",
        Err(AlignmentFailure::SoftFailure {
            reason: SoftFailureReason::Timeout(_),
        }) => "timeout",
        Err(AlignmentFailure::SoftFailure { .. } | AlignmentFailure::Error { .. }) => "error",
    }
}

fn duration_ms(duration: Duration) -> f64 {
    duration.as_secs_f64() * 1000.0
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use generic_a_star::cost::AStarCost as _;
    use lib_tsalign::a_star_aligner::{
        alignment_geometry::AlignmentRange,
        alignment_result::alignment::Alignment,
        template_switch_distance::{
            EqualCostRange, TemplateSwitchAncestor, TemplateSwitchDescendant,
            TemplateSwitchDirection,
        },
    };
    use lib_tsalign::costs::U64Cost;

    use super::*;
    use crate::common::{
        aligner::{fpa::FpaAlignmentStatistics, result::TwitcherAlignmentWithStatistics},
        contig::ContigName,
        coords::GenomePosition,
    };

    #[test]
    fn parse_proc_status_reads_peak_fields() {
        let status =
            "Name:\ttwitcher\nVmPeak:\t  123456 kB\nVmSize:\t   1000 kB\nVmHWM:\t     78 kB\n";
        assert_eq!(
            parse_proc_status(status),
            MemoryUsage {
                peak_rss: Some(78 * 1024),
                peak_vsz: Some(123_456 * 1024),
            }
        );
    }

    /// A status file without the peak fields (or a garbled one) must not lose the other values.
    #[test]
    fn parse_proc_status_tolerates_missing_and_broken_fields() {
        assert_eq!(parse_proc_status(""), MemoryUsage::default());
        assert_eq!(parse_proc_status("no colon here"), MemoryUsage::default());
        assert_eq!(
            parse_proc_status("VmPeak:\tnonsense\nVmHWM:\t8 kB\n"),
            MemoryUsage {
                peak_rss: Some(8 * 1024),
                peak_vsz: None,
            }
        );
    }

    #[test]
    fn parse_kb_accepts_values_with_and_without_unit() {
        assert_eq!(parse_kb("  12 kB"), Some(12 * 1024));
        assert_eq!(parse_kb("12"), Some(12 * 1024));
        assert_eq!(parse_kb("-1 kB"), None);
    }

    /// The peak memory of the test process itself must be readable and non-zero on Linux.
    #[test]
    fn current_process_memory_is_observable() {
        let memory = MemoryUsage::of_current_process();
        assert!(memory.peak_rss.is_some_and(|rss| rss > 0));
        assert!(memory.peak_vsz.is_some_and(|vsz| vsz > 0));
    }

    fn region() -> GenomeRegion {
        GenomeRegion::from_incl_incl(
            GenomePosition::new_0(ContigName::new(b"chr1"), 100),
            Some(GenomePosition::new_0(ContigName::new(b"chr1"), 120)),
        )
        .unwrap()
    }

    fn exec() -> ExecStats {
        ExecStats {
            wall: Duration::from_millis(1500),
            memory: MemoryUsage {
                peak_rss: Some(2048),
                peak_vsz: Some(4096),
            },
            memory_limit: 8192,
        }
    }

    /// Build an FPA success whose alignment contains `ts_count` template switches.
    fn success(ts_count: usize) -> TwitcherAlignment {
        let mut alignment = Alignment::new();
        for _ in 0..ts_count {
            alignment.push_n(
                1,
                AlignmentType::TemplateSwitchEntrance {
                    first_offset: 0,
                    equal_cost_range: EqualCostRange::new_invalid(),
                    descendant: TemplateSwitchDescendant::Reference,
                    ancestor: TemplateSwitchAncestor::Reference,
                    direction: TemplateSwitchDirection::Reverse,
                },
            );
            alignment.push_n(1, AlignmentType::SecondaryMatch);
            alignment.push_n(
                1,
                AlignmentType::TemplateSwitchExit {
                    anti_descendant_gap: 0,
                },
            );
        }
        alignment.push_n(5, AlignmentType::PrimaryMatch);

        TwitcherAlignment::new(
            Ok(TwitcherAlignmentWithStatistics::new(
                alignment,
                U64Cost::from_primitive(42),
                FpaAlignmentStatistics {
                    duration: Duration::from_millis(250),
                    ranges: AlignmentRange::new_complete(20, 20),
                }
                .into(),
            )),
            exec(),
        )
    }

    /// A `Write` that keeps the written bytes readable after the writer was dropped.
    #[derive(Clone, Default)]
    struct SharedBuffer(Arc<Mutex<Vec<u8>>>);

    impl Write for SharedBuffer {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            self.0.lock().unwrap().write(buf)
        }

        fn flush(&mut self) -> std::io::Result<()> {
            Ok(())
        }
    }

    fn write_rows(alignments: &[TwitcherAlignment]) -> String {
        let buffer = SharedBuffer::default();
        {
            let writer = AlignmentStatsWriter::new(Box::new(buffer.clone()), "fpa");
            for alignment in alignments {
                writer.write(
                    alignment,
                    &AlignmentStatsContext {
                        cluster_region: &region(),
                        source: AlignmentSource::Computed,
                        reference_length: 200,
                        query_length: 210,
                    },
                );
            }
        }
        let bytes = buffer.0.lock().unwrap().clone();
        String::from_utf8(bytes).unwrap()
    }

    /// The header is written once, and a successful alignment fills in the columns it can.
    #[test]
    fn successful_alignment_row() {
        let rows = write_rows(&[success(2)]);
        let mut lines = rows.lines();
        assert_eq!(
            lines.next().unwrap(),
            "cluster_region,source,aligner,outcome,cost,ts_num,wall_ms,align_ms,peak_rss,peak_vsz,\
             memory_limit,opened_nodes,closed_nodes,suboptimal_nodes,suboptimal_ratio,\
             cost_per_base,reference_length,query_length,error"
        );
        assert_eq!(
            lines.next().unwrap(),
            "chr1:101-121,computed,fpa,ts,42,2,1500.0,250.0,2048,4096,8192,,,,,,200,210,"
        );
        assert!(lines.next().is_none());
    }

    /// An alignment without a template switch is reported as such, and still carries its cost.
    #[test]
    fn alignment_without_template_switch_is_no_ts() {
        let rows = write_rows(&[success(0)]);
        let row = rows.lines().nth(1).unwrap();
        assert!(row.contains(",fpa,no_ts,42,0,"), "{row}");
    }

    /// Failures have no aligner statistics, but must still report runtime, memory and a reason.
    #[test]
    fn failure_rows_keep_execution_statistics() {
        let failures = [
            AlignmentFailure::oom(),
            AlignmentFailure::timeout(Duration::from_secs(1)),
            AlignmentFailure::error("boom"),
        ];
        let alignments: Vec<_> = failures
            .into_iter()
            .map(|failure| TwitcherAlignment::new(Err(failure), exec()))
            .collect();

        let rows = write_rows(&alignments);
        let rows: Vec<_> = rows.lines().skip(1).collect();
        assert_eq!(
            rows,
            [
                "chr1:101-121,computed,fpa,oom,,,1500.0,,2048,4096,8192,,,,,,200,210,",
                "chr1:101-121,computed,fpa,timeout,,,1500.0,,2048,4096,8192,,,,,,200,210,",
                "chr1:101-121,computed,fpa,error,,,1500.0,,2048,4096,8192,,,,,,200,210,boom",
            ]
        );
    }
}