nornir 0.2.0

//! Bench API — the contract child workspaces (znippy, holger, ...)
//! implement so a `nornir release` can drive their benches.
//!
//! Design contract (intentionally tiny):
//!
//! 1. Each repo adds `nornir = { path = "../nornir" }` (or git/version)
//!    under `[dev-dependencies]` ONLY — nothing about nornir leaks into
//!    the published crate.
//! 2. Each repo provides one `examples/nornir-bench.rs` of literally:
//!    ```ignore
//!    fn main() -> anyhow::Result<()> {
//!        nornir::bench::api::run_main_json()
//!    }
//!    ```
//!    Examples can use dev-deps; bins cannot. That's why this is an
//!    example, not a bin.
//! 3. Each bench is a tiny struct implementing [`Bencher`] and
//!    registered via [`register_bench!`]. `inventory` collects them at
//!    link time — no manual registry list to maintain.
//! 4. `cargo run --release --example nornir-bench` prints a
//!    [`crate::bench::BenchRun`] as JSON to stdout. Nornir's release
//!    pipeline subprocess-spawns that, parses stdout, writes the run
//!    into Urðr tagged with the release_id + per-repo git SHA + dep
//!    graph snapshot id.
//!
//! Why on stdout and not over a socket/pipe protocol? Same reason as
//! `cargo metadata`: stdout-JSON is the simplest possible contract
//! between two unrelated build trees. Nornir is the orchestrator;
//! children stay passive and `cargo run`-able in isolation.

use anyhow::Result;
use chrono::Utc;

use super::{BenchResult, BenchRun};

/// A single bench. `id` is logged into Urðr verbatim; convention is
/// `<repo>.<scenario>` (e.g. `"holger.artifact_throughput_st"`).
///
/// `run` returns the raw [`BenchResult`]; the runner stamps the run
/// envelope (date, version, machine, cores) automatically.
pub trait Bencher: Sync {
    fn id(&self) -> &'static str;
    fn run(&self) -> Result<BenchResult>;
}

/// Registers a `&'static dyn Bencher` so the runner discovers it at
/// link time via `inventory`.
///
/// Usage:
/// ```ignore
/// struct ArtifactThroughput;
/// impl nornir::bench::api::Bencher for ArtifactThroughput { /* … */ }
/// nornir::bench::register_bench!(ArtifactThroughput);
/// ```
/// Registers a bencher instance. The given expression is evaluated
/// (and leaked into `'static`) the first time the runner iterates the
/// registry. Use a zero-sized unit struct + `register_bench!(Foo)` for
/// the common case, or any constructor expression for stateful ones.
///
/// Order defaults to `0`. When stages of a pipeline must run in a
/// deterministic sequence (output of stage N feeds stage N+1) use
/// [`register_bench_ordered!`] with ascending integers; the runner
/// sorts by `order` before executing, breaking ties on `id()`.
///
/// Usage:
/// ```ignore
/// struct ArtifactThroughput;
/// impl nornir::bench::api::Bencher for ArtifactThroughput { /* … */ }
/// nornir::register_bench!(ArtifactThroughput);
///
/// // or with a constructor:
/// nornir::register_bench!(ArtifactThroughput::with_size(1 << 20));
/// ```
#[macro_export]
macro_rules! register_bench {
    ($expr:expr) => {
        $crate::bench::api::inventory_submit! {
            $crate::bench::api::BencherRegistration {
                order: 0,
                make: || {
                    let b: ::std::boxed::Box<dyn $crate::bench::api::Bencher> =
                        ::std::boxed::Box::new($expr);
                    ::std::boxed::Box::leak(b)
                },
            }
        }
    };
}

/// Like [`register_bench!`] but pins an explicit execution order.
/// Lower `$order` runs first; equal orders fall back to `id()`
/// lexicographic order. Use for strict pipelines (e.g.
/// `bz2 → pbf` must finish before `pbf → geoparquet` starts).
///
/// Usage:
/// ```ignore
/// nornir::register_bench_ordered!(0, Stage1BzToPbf);
/// nornir::register_bench_ordered!(1, Stage2PbfToGeo);
/// ```
#[macro_export]
macro_rules! register_bench_ordered {
    ($order:expr, $expr:expr) => {
        $crate::bench::api::inventory_submit! {
            $crate::bench::api::BencherRegistration {
                order: $order,
                make: || {
                    let b: ::std::boxed::Box<dyn $crate::bench::api::Bencher> =
                        ::std::boxed::Box::new($expr);
                    ::std::boxed::Box::leak(b)
                },
            }
        }
    };
}

// Re-export so users only need `nornir::register_bench!` plus the
// inventory submission target below.
pub use inventory::submit as inventory_submit;

/// Inventory entry. The function pointer dodges issues with const
/// `&dyn Trait` references — each registration just returns its
/// `'static` reference on demand.
///
/// `order` is a stable sort key applied before execution. Lower runs
/// first; ties broken on `Bencher::id()`. Default 0 (use
/// [`register_bench_ordered!`] to pin pipeline stages).
pub struct BencherRegistration {
    pub order: i32,
    pub make: fn() -> &'static dyn Bencher,
}

inventory::collect!(BencherRegistration);

/// The one-liner each repo's `examples/nornir-bench.rs` calls.
///
/// Collects every `register_bench!`'d entry, sorts by `(order, id)`
/// so pipeline stages run in a deterministic sequence, and executes
/// them sequentially (no parallel scheduling — bench machines should
/// be quiescent).
///
/// Builds a full [`BenchRun`] envelope and prints it as one JSON line
/// on stdout. Errors from individual benches are converted to red
/// [`super::TestOutcome`] entries so a bad bench degrades to a
/// recorded failure rather than killing the whole run — **except**
/// when `NORNIR_BENCH_STOP_ON_ERROR=1`, in which case any failure
/// aborts the remaining stages (use for pipelines where stage N+1
/// reads stage N's artifact and would fail anyway).
pub fn run_main_json() -> Result<()> {
    let mut regs: Vec<&'static BencherRegistration> =
        inventory::iter::<BencherRegistration>().collect();
    // Resolve `make` once so we can sort on id() too.
    let mut resolved: Vec<(i32, &'static dyn Bencher)> =
        regs.drain(..).map(|r| (r.order, (r.make)())).collect();
    resolved.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.id().cmp(b.1.id())));

    let stop_on_error = std::env::var("NORNIR_BENCH_STOP_ON_ERROR")
        .map(|v| v != "0" && !v.is_empty())
        .unwrap_or(false);

    let mut results: Vec<BenchResult> = Vec::new();
    let mut tests: Vec<super::TestOutcome> = Vec::new();
    for (_order, b) in resolved {
        let id = b.id().to_string();
        let start = std::time::Instant::now();
        match b.run() {
            Ok(r) => {
                results.push(r);
                tests.push(super::TestOutcome {
                    name: id,
                    passed: true,
                    duration_ms: Some(start.elapsed().as_secs_f64() * 1000.0),
                    message: None,
                });
            }
            Err(e) => {
                tests.push(super::TestOutcome {
                    name: id,
                    passed: false,
                    duration_ms: Some(start.elapsed().as_secs_f64() * 1000.0),
                    message: Some(format!("{e:#}")),
                });
                if stop_on_error {
                    break;
                }
            }
        }
    }
    let now = Utc::now();
    let run = BenchRun {
        date: now.format("%Y-%m-%d").to_string(),
        timestamp: Some(now.to_rfc3339()),
        // The bench runner is often a separate crate (e.g. a workspace's
        // `xtask`), so its `CARGO_PKG_VERSION` is not the version under test.
        // `NORNIR_BENCH_VERSION` lets the caller stamp the *subject* repo's
        // version (e.g. `nornir bench run` / a CI step sets it); fall back to
        // the bench crate's own version.
        version: std::env::var("NORNIR_BENCH_VERSION")
            .ok()
            .filter(|v| !v.is_empty())
            .unwrap_or_else(|| env!("CARGO_PKG_VERSION").to_string()),
        machine: std::env::var("NORNIR_MACHINE").unwrap_or_default(),
        cores: num_cpus_best_effort() as u32,
        results,
        tests,
    };
    println!("{}", serde_json::to_string(&run)?);
    Ok(())
}

fn num_cpus_best_effort() -> usize {
    std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1)
}

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

    struct Demo;
    impl Bencher for Demo {
        fn id(&self) -> &'static str { "demo.always_42" }
        fn run(&self) -> Result<BenchResult> {
            let mut metrics = serde_json::Map::new();
            metrics.insert("answer".into(), serde_json::json!(42));
            Ok(BenchResult { name: "demo".into(), metrics })
        }
    }
    crate::register_bench!(Demo);

    #[test]
    fn registry_includes_demo() {
        let ids: Vec<&'static str> =
            inventory::iter::<BencherRegistration>().map(|r| (r.make)().id()).collect();
        assert!(ids.contains(&"demo.always_42"), "registry missing demo: {ids:?}");
    }
}