Skip to main content

rule_benchmark/
rule_benchmark.rs

1use std::fs;
2use std::hint::black_box;
3use std::time::{Duration, Instant};
4
5use tempfile::TempDir;
6use yarp_cli::reducers::{StreamReducer, configured_memory_bound};
7use yarp_cli::rules::{PackRequest, Registry, Selection};
8use yarp_rule_pack::{SourcePack, compile};
9
10const EXTERNAL_RULES: usize = 1_000;
11const MATCH_SAMPLES: usize = 200;
12const BUILTIN_TARGET_US: u128 = 100;
13const EXTERNAL_TARGET_US: u128 = 5_000;
14const STREAM_TARGET_MEGABYTES_PER_SECOND: f64 = 100.0;
15
16fn main() -> Result<(), String> {
17    let builtin = benchmark_builtin()?;
18    let (external_cold, external, pack_bytes) = benchmark_external()?;
19    let (throughput, memory_bound) = benchmark_stream()?;
20    println!("built_in_match_p95_us: {}", builtin.as_micros());
21    println!(
22        "external_1000_first_open_match_us: {}",
23        external_cold.as_micros()
24    );
25    println!("external_1000_open_match_p95_us: {}", external.as_micros());
26    println!("external_pack_bytes: {pack_bytes}");
27    println!("stream_megabytes_per_second: {throughput:.2}");
28    println!("configured_stream_memory_bound_bytes: {memory_bound}");
29    if builtin.as_micros() > BUILTIN_TARGET_US {
30        return Err(format!(
31            "built-in p95 exceeded {BUILTIN_TARGET_US} microseconds"
32        ));
33    }
34    if external.as_micros() > EXTERNAL_TARGET_US {
35        return Err(format!(
36            "external pack p95 exceeded {EXTERNAL_TARGET_US} microseconds"
37        ));
38    }
39    if throughput < STREAM_TARGET_MEGABYTES_PER_SECOND {
40        return Err(format!(
41            "stream throughput fell below {STREAM_TARGET_MEGABYTES_PER_SECOND:.0} MB/s"
42        ));
43    }
44    Ok(())
45}
46
47fn benchmark_builtin() -> Result<Duration, String> {
48    let arguments = vec![
49        "cargo".to_owned(),
50        "test".to_owned(),
51        "--workspace".to_owned(),
52    ];
53    let mut samples = Vec::with_capacity(MATCH_SAMPLES);
54    for _ in 0..MATCH_SAMPLES {
55        let mut registry = Registry::builtins_only();
56        let started = Instant::now();
57        let selection = registry.select(black_box(&arguments))?;
58        if !matches!(selection, Selection::Reduce(_)) {
59            return Err("built-in benchmark command was not selected".to_owned());
60        }
61        samples.push(started.elapsed());
62    }
63    Ok(p95(samples))
64}
65
66fn benchmark_external() -> Result<(Duration, Duration, usize), String> {
67    let directory = TempDir::new().map_err(|error| error.to_string())?;
68    let source_root = directory.path().join("source");
69    let rules_root = source_root.join("rules");
70    fs::create_dir_all(&rules_root).map_err(|error| error.to_string())?;
71    let mut paths = Vec::with_capacity(EXTERNAL_RULES);
72    for index in 0..EXTERNAL_RULES {
73        let relative = format!("rules/rule-{index:04}.json");
74        let id = format!("bench/rule-{index:04}");
75        let program = format!("bench{index:04}");
76        let body = format!(
77            "{{\"id\":{id:?},\"match\":{{\"program\":[{program:?}]}},\"action\":\"reduce\",\"reducer\":{{\"kind\":\"list_summary\"}},\"success\":{{\"max_line_bytes\":16384,\"max_output_bytes\":32768,\"min_savings_bytes\":120,\"min_savings_basis_points\":1000}},\"failure\":{{\"max_line_bytes\":16384,\"max_output_bytes\":65536,\"min_savings_bytes\":120,\"min_savings_basis_points\":500}}}}"
78        );
79        fs::write(source_root.join(&relative), body).map_err(|error| error.to_string())?;
80        paths.push(relative);
81    }
82    let manifest = serde_json::json!({
83        "schema_version": 1,
84        "id": "benchmark-pack",
85        "rules": paths,
86    });
87    fs::write(
88        source_root.join("pack.json"),
89        serde_json::to_vec(&manifest).map_err(|error| error.to_string())?,
90    )
91    .map_err(|error| error.to_string())?;
92    let source = SourcePack::load(&source_root)?;
93    let body = compile(&source)?;
94    let pack_path = directory.path().join("benchmark.yrp");
95    fs::write(&pack_path, &body).map_err(|error| error.to_string())?;
96    let arguments = vec!["bench0999".to_owned()];
97    let request = PackRequest {
98        path: pack_path,
99        expected_digest: Some(source.source_digest),
100        expected_compiled_digest: None,
101    };
102    let cold_started = Instant::now();
103    let mut cold_registry = Registry::load(std::slice::from_ref(&request))?;
104    if !matches!(
105        cold_registry.select(black_box(&arguments))?,
106        Selection::Reduce(_)
107    ) {
108        return Err("external benchmark command was not selected".to_owned());
109    }
110    let cold = cold_started.elapsed();
111    let mut samples = Vec::with_capacity(MATCH_SAMPLES);
112    for _ in 0..MATCH_SAMPLES {
113        let started = Instant::now();
114        let mut registry = Registry::load(std::slice::from_ref(&request))?;
115        let selection = registry.select(black_box(&arguments))?;
116        if !matches!(selection, Selection::Reduce(_)) {
117            return Err("external benchmark command was not selected".to_owned());
118        }
119        samples.push(started.elapsed());
120    }
121    Ok((cold, p95(samples), body.len()))
122}
123
124fn benchmark_stream() -> Result<(f64, usize), String> {
125    let mut registry = Registry::builtins_only();
126    let Selection::Reduce(selected) = registry.select(&["rg".to_owned(), "needle".to_owned()])?
127    else {
128        return Err("stream benchmark rule was not selected".to_owned());
129    };
130    let chunk = b"src/file.rs:42:matching line with useful context\n".repeat(160);
131    let bytes = 1_024 * 1_024 * 1_024;
132    let started = Instant::now();
133    let mut reducer = StreamReducer::new(&selected.rule)?;
134    let mut processed = 0;
135    while processed < bytes {
136        reducer.push(black_box(&chunk));
137        processed += chunk.len();
138    }
139    black_box(reducer.finish(true, None));
140    let processed = u32::try_from(processed)
141        .map_err(|_| "stream benchmark byte count does not fit u32".to_owned())?;
142    let throughput = f64::from(processed) / 1_000_000.0 / started.elapsed().as_secs_f64();
143    Ok((throughput, configured_memory_bound(&selected.rule)?))
144}
145
146fn p95(mut samples: Vec<Duration>) -> Duration {
147    samples.sort_unstable();
148    let index = samples
149        .len()
150        .saturating_mul(95)
151        .div_ceil(100)
152        .saturating_sub(1);
153    samples[index]
154}