nordisk/nornir

crates/mimir-room/examples/room_rss.rs

//! What a room costs in memory: open a lake, build the current room (every
//! consumed repo, ring off) and print the resident set before and after.
//!
//! ```text
//! cargo run --release -p nornir-mimir-room --example room_rss -- <builds>/mimir_v2/lake [ring|cur] [repo]
//! ```

fn rss_mb() -> u64 {
    status_mb("VmRSS:")
}

fn status_mb(field: &str) -> u64 {
    let s = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
    s.lines().find(|l| l.starts_with(field)).and_then(|l| l.split_whitespace().nth(1)).and_then(|v| v.parse::<u64>().ok()).unwrap_or(0) / 1024
}

fn main() -> anyhow::Result<()> {
    let lake = std::env::args().nth(1).ok_or_else(|| anyhow::anyhow!("usage: room_rss <lake dir>"))?;
    // `repair`: open the catalog alone and close it (a copy of a lake taken
    // while its server wrote is repaired on open) — nothing else
    if std::env::args().nth(2).as_deref() == Some("repair") {
        let rt = tokio::runtime::Builder::new_current_thread().enable_all().build()?;
        drop(rt.block_on(skade::open(&lake))?);
        println!("{lake}: catalog opened and closed");
        return Ok(());
    }
    let log = mimir_ingest::EventLog::open_lake_blocking(std::path::Path::new(&lake))?;
    // `cols`: what a ring row costs, per column — the ring rows of the first
    // 200 ring owners (by crate node), every column of the node and edge tables
    if std::env::args().nth(2).as_deref() == Some("cols") {
        use skade::arrow_array::Array;
        let crates = log.block_on(log.scan_rows_where(mimir_ingest::Kind::Nodes, &skade::ScanFilter::is_in("kind", ["crate"]), &["kind"]))?;
        let mut owners = std::collections::BTreeSet::new();
        for b in &crates.batches {
            let (repo, krate) = (b.column_by_name("repo").unwrap(), b.column_by_name("krate").unwrap());
            let (repo, krate) = (repo.as_any().downcast_ref::<skade::arrow_array::StringArray>().unwrap(), krate.as_any().downcast_ref::<skade::arrow_array::StringArray>().unwrap());
            for i in 0..b.num_rows() {
                if repo.value(i) == "@universe" {
                    owners.insert(krate.value(i).to_string());
                }
            }
        }
        let some: Vec<&str> = owners.iter().take(200).map(String::as_str).collect();
        for kind in [mimir_ingest::Kind::Nodes, mimir_ingest::Kind::Edges] {
            let all: Vec<String> = log.spec(kind).schema.fields().iter().map(|f| f.name().clone()).collect();
            let all: Vec<&str> = all.iter().map(String::as_str).collect();
            let rows = log.block_on(log.scan_rows_where(kind, &skade::ScanFilter::is_in("krate", some.iter().copied()), &all))?;
            let n: usize = rows.batches.iter().map(|b| b.num_rows()).sum();
            let defined = rows.batches.iter().filter_map(|b| b.column_by_name("defined_here")).map(|c| c.as_any().downcast_ref::<skade::arrow_array::BooleanArray>().map_or(0, |a| a.true_count())).sum::<usize>();
            println!("{kind:?}: {n} rows of {} owners ({} of {} ring owners), defined_here {defined}", some.len(), some.len(), owners.len());
            let schema = rows.batches[0].schema();
            for (i, f) in schema.fields().iter().enumerate() {
                let bytes: usize = rows.batches.iter().map(|b| b.column(i).get_array_memory_size()).sum();
                println!("  {:<20} {:>8.1} B/row", f.name(), bytes as f64 / n.max(1) as f64);
            }
        }
        return Ok(());
    }
    let r0 = rss_mb();
    let t = std::time::Instant::now();
    // a second argument `ring`: the dependency ring on (the panes' ring view)
    let ring = std::env::args().nth(2).as_deref() == Some("ring");
    // a third: one repo's room (its ring view with `ring`), as a pane asks
    let base = match std::env::args().nth(3) {
        Some(repo) => mimir_room::Scope::repo(&repo, None),
        None => mimir_room::Scope::current(),
    };
    let scope = base.with_universe(ring);
    let opened = log.block_on(mimir_room::Room::open(&log, &scope))?;
    let built = rss_mb();
    println!("room: {} nodes, {} ms · RSS {r0} MB → {built} MB, peak {} MB", opened.room.len(), t.elapsed().as_millis(), status_mb("VmHWM:"));
    drop(opened);
    mimir_room::Room::clear_cache();
    mimir_room::Room::clear_ring_cache();
    let gone = rss_mb();
    println!("room dropped: RSS {gone} MB — the room held {} MB", built.saturating_sub(gone));
    drop(log);
    println!("log dropped: RSS {} MB", rss_mb());
    Ok(())
}