crates/nornir-flight/src/bench.rs
//! **Bench the transport hot path** (feature `bench`) — `do_put` N `RecordBatch`es
//! into the seam over a real localhost Flight connection, then `do_get` them
//! back, measuring put and get throughput (MiB/s and rows/s). This is the seam's
//! actual work: Arrow IPC framing on/off the wire through the tonic stream.
//!
//! Registered with the shared `nornir-bench` harness via [`register_bench!`]; the
//! harness stamps host cores + peak memory telemetry onto the row automatically.
//! OFF by default — a release/library build carries no `nornir-bench` dependency.
use std::net::TcpListener;
use std::sync::Arc;
use std::time::{Duration, Instant};
use arrow_array::{Int64Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema};
use nornir_bench::api::Bencher;
use nornir_bench::BenchResult;
use serde_json::json;
use crate::{FlightConn, FlightRoutes, FlightServer, MemFlightStore};
/// The put/get throughput bencher: `batches` × `rows_per_batch` two-column
/// (Int64 + Utf8) records, pushed then pulled over a localhost Flight seam. The
/// data is generated deterministically (no clock, no rng), so the number is
/// stable and comparable across hosts.
pub struct FlightThroughput {
rows_per_batch: usize,
batches: usize,
}
impl Default for FlightThroughput {
fn default() -> Self {
// ~1M rows total — enough to dwarf connection-setup cost.
//
// PREVIEW (`NORNIR_BENCH_PREVIEW=1`): the SAME put/get path over a real
// localhost Flight connection, same metric keys, 1/50 the rows. It exists so a
// bencher (and its whole persist → render chain) can be verified without waiting
// out a real pass. The harness stamps a preview run onto `<host>-preview`, so its
// number can never land in — or overwrite — the real host's README row.
if nornir_bench::api::preview_run() {
Self { rows_per_batch: 1_000, batches: 20 }
} else {
Self { rows_per_batch: 50_000, batches: 20 }
}
}
}
impl FlightThroughput {
/// A bencher over `batches` batches of `rows_per_batch` rows each.
pub fn with_size(rows_per_batch: usize, batches: usize) -> Self {
Self { rows_per_batch, batches }
}
}
/// `batches` deterministic two-column record batches.
fn make_batches(rows_per_batch: usize, batches: usize) -> Vec<RecordBatch> {
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int64, false),
Field::new("name", DataType::Utf8, true),
]));
(0..batches)
.map(|b| {
let ids: Vec<i64> = (0..rows_per_batch).map(|r| (b * rows_per_batch + r) as i64).collect();
let names: Vec<Option<String>> = (0..rows_per_batch)
.map(|r| if r % 7 == 0 { None } else { Some(format!("row-{b}-{r}")) })
.collect();
RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int64Array::from(ids)), Arc::new(StringArray::from(names))],
)
.expect("bench batch builds")
})
.collect()
}
impl Bencher for FlightThroughput {
fn id(&self) -> &'static str {
"nornir-flight.put_get_throughput"
}
fn run(&self) -> anyhow::Result<BenchResult> {
let rows_per_batch = self.rows_per_batch;
let n_batches = self.batches;
let rt = tokio::runtime::Builder::new_multi_thread().enable_all().build()?;
rt.block_on(async move {
let batches = make_batches(rows_per_batch, n_batches);
let bytes: usize = batches.iter().map(|b| b.get_array_memory_size()).sum();
let total_rows = (rows_per_batch * n_batches) as f64;
let mib = bytes as f64 / (1024.0 * 1024.0);
// Start an open (no-auth) server on a free localhost port.
let addr = TcpListener::bind("127.0.0.1:0")?.local_addr()?;
let routes = FlightRoutes::shared(Arc::new(MemFlightStore::new()));
// A32-14, stays a task: the bench's local Flight server (feature bench), never in a box.
tokio::spawn(async move {
let _ = FlightServer::serve(routes, addr).await;
});
for _ in 0..300 {
if tokio::net::TcpStream::connect(addr).await.is_ok() {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
let mut c = FlightConn::connect(format!("http://{addr}"))
.await
.map_err(|e| anyhow::anyhow!(e.to_string()))?;
let t_put = Instant::now();
c.put("bench", batches.clone())
.await
.map_err(|e| anyhow::anyhow!(e.to_string()))?;
let put_s = t_put.elapsed().as_secs_f64().max(1e-9);
let t_get = Instant::now();
let got = c.get("bench").await.map_err(|e| anyhow::anyhow!(e.to_string()))?;
let get_s = t_get.elapsed().as_secs_f64().max(1e-9);
// The Flight encoder may re-chunk large batches, so compare total
// rows (which are preserved) rather than the batch count.
let got_rows: usize = got.iter().map(|b| b.num_rows()).sum();
anyhow::ensure!(
got_rows == rows_per_batch * n_batches,
"round-trip row count mismatch: got {got_rows}",
);
let mut m = serde_json::Map::new();
m.insert("batches".into(), json!(n_batches as f64));
m.insert("rows".into(), json!(total_rows));
m.insert("mib".into(), json!(mib));
m.insert("put_mib_per_sec".into(), json!(mib / put_s));
m.insert("get_mib_per_sec".into(), json!(mib / get_s));
m.insert("put_rows_per_sec".into(), json!(total_rows / put_s));
m.insert("get_rows_per_sec".into(), json!(total_rows / get_s));
Ok(BenchResult::measured(self.id(), m))
})
}
}
// Register with the shared harness so `nornir bench` / the example entrypoint
// picks up the Flight transport alongside every other nordisk bencher.
nornir_bench::register_bench!(FlightThroughput::default());
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test(flavor = "multi_thread")]
async fn throughput_bencher_reports_positive_put_get() {
// Small size so the unit test is quick; the real run uses the default.
let b = FlightThroughput::with_size(2_000, 3);
assert_eq!(b.id(), "nornir-flight.put_get_throughput");
// run() drives its own runtime; call it on a blocking thread so it does
// not nest inside this test's runtime.
let r = tokio::task::spawn_blocking(move || b.run()).await.unwrap().expect("bench runs");
let put = r.metrics.get("put_mib_per_sec").and_then(|v| v.as_f64()).unwrap();
let get = r.metrics.get("get_mib_per_sec").and_then(|v| v.as_f64()).unwrap();
assert!(put > 0.0, "put throughput must be positive");
assert!(get > 0.0, "get throughput must be positive");
}
}