indexed_output_benchmark/
indexed_output_benchmark.rs

1use std::fs;
2use std::io::Write as _;
3use std::path::{Path, PathBuf};
4use std::process::{Command, Stdio};
5use std::time::{Duration, Instant};
6
7use serde_json::json;
8use tempfile::TempDir;
9use yarp_cli::archive::{Archive, CallIdentity, SessionIdentity, SourceCompleteness};
10
11const WARMUP: usize = 10;
12const SAMPLES: usize = 100;
13const TARGET: Duration = Duration::from_millis(20);
14const PARSER_TARGET: Duration = Duration::from_millis(1);
15
16fn main() -> Result<(), String> {
17    let executable = release_yarp()?;
18    let directory = TempDir::new().map_err(|error| error.to_string())?;
19    let database = directory.path().join("archive/tool-calls.sqlite3");
20    let config_home = directory.path().join("config");
21    fs::create_dir_all(config_home.join("yarp")).map_err(|error| error.to_string())?;
22    let database_text = serde_json::to_string(
23        database
24            .to_str()
25            .ok_or_else(|| "benchmark archive path is not UTF-8".to_owned())?,
26    )
27    .map_err(|error| error.to_string())?;
28    fs::write(
29        config_home.join("yarp/config.toml"),
30        format!("version = 1\n[archive]\npath = {database_text}\n"),
31    )
32    .map_err(|error| error.to_string())?;
33    let mut archive = Archive::open_path(database.clone())?;
34    let session = SessionIdentity {
35        agent: "benchmark".to_owned(),
36        account: "local".to_owned(),
37        source_session_id: "indexed-output".to_owned(),
38        started_at_ms: Some(1),
39    };
40    let call = CallIdentity {
41        source_call_id: "query".to_owned(),
42        tool_name: "exec_command".to_owned(),
43        provider: None,
44        model: None,
45        working_directory: None,
46        started_at_ms: 2,
47        requires_streams: false,
48    };
49    let archive_ref = archive.begin_call(&session, &call, &json!({}), &json!({}), 2)?;
50    let source = "path/to/file.rs:42: needle repeated output for bounded query\n".repeat(17_000);
51    archive.result_text(&session, "query", &source, SourceCompleteness::Complete, 3)?;
52    drop(archive);
53
54    let search = benchmark_command(
55        &executable,
56        &config_home,
57        &["search", &archive_ref, "needle", "--max-results", "20"],
58    )?;
59    let read = benchmark_command(
60        &executable,
61        &config_home,
62        &["read", &archive_ref, "result_text", "1:200"],
63    )?;
64    let result = benchmark_result_reducer(&executable, &config_home, &archive_ref)?;
65    let representative_source = "set -o pipefail && rg TODO . | sort | uniq | head -50";
66    let representative_parser = benchmark_parser(representative_source, WARMUP, SAMPLES)?;
67    let maximum_source = format!("rg {}", "x".repeat(256 * 1024 - 3));
68    let maximum_parser = benchmark_parser(&maximum_source, 2, 10)?;
69    print_stats("search_1m", &search);
70    print_stats("read_12k", &read);
71    print_stats("result_reducer_16k", &result);
72    print_stats("parser_representative", &representative_parser);
73    print_stats("parser_maximum_accepted", &maximum_parser);
74    println!(
75        "parser_representative_input_bytes: {}",
76        representative_source.len()
77    );
78    println!(
79        "parser_maximum_accepted_input_bytes: {}",
80        maximum_source.len()
81    );
82    for (name, samples) in [
83        ("search", &search),
84        ("read", &read),
85        ("result reducer", &result),
86    ] {
87        if percentile(samples, 95) > TARGET {
88            return Err(format!("{name} p95 exceeded {} ms", TARGET.as_millis()));
89        }
90    }
91    if percentile(&representative_parser, 95) > PARSER_TARGET {
92        return Err(format!(
93            "representative parser p95 exceeded {} ms",
94            PARSER_TARGET.as_millis()
95        ));
96    }
97    Ok(())
98}
99
100fn benchmark_parser(
101    source: &str,
102    warmup: usize,
103    sample_count: usize,
104) -> Result<Vec<Duration>, String> {
105    let mut samples = Vec::with_capacity(sample_count);
106    for iteration in 0..warmup.saturating_add(sample_count) {
107        let started = Instant::now();
108        std::hint::black_box(yarp_cli::rewrite::select_result_plan(std::hint::black_box(
109            source,
110        )))
111        .map_err(|error| format!("parser benchmark command was rejected: {error}"))?;
112        let elapsed = started.elapsed();
113        if iteration >= warmup {
114            samples.push(elapsed);
115        }
116    }
117    Ok(samples)
118}
119
120fn release_yarp() -> Result<PathBuf, String> {
121    let current = std::env::current_exe().map_err(|error| error.to_string())?;
122    let release = current
123        .parent()
124        .and_then(Path::parent)
125        .ok_or_else(|| "could not find release directory".to_owned())?;
126    let executable = release.join(if cfg!(windows) { "yarp.exe" } else { "yarp" });
127    if !executable.is_file() {
128        return Err(format!(
129            "{} is missing; run cargo build --release --workspace first",
130            executable.display()
131        ));
132    }
133    Ok(executable)
134}
135
136fn benchmark_command(
137    executable: &Path,
138    config_home: &Path,
139    arguments: &[&str],
140) -> Result<Vec<Duration>, String> {
141    let mut samples = Vec::with_capacity(SAMPLES);
142    for iteration in 0..WARMUP + SAMPLES {
143        let started = Instant::now();
144        let output = Command::new(executable)
145            .args(arguments)
146            .env("XDG_CONFIG_HOME", config_home)
147            .output()
148            .map_err(|error| error.to_string())?;
149        let elapsed = started.elapsed();
150        if !output.status.success() {
151            return Err(format!(
152                "{} failed: {}",
153                arguments[0],
154                String::from_utf8_lossy(&output.stderr)
155            ));
156        }
157        if iteration >= WARMUP {
158            samples.push(elapsed);
159        }
160    }
161    Ok(samples)
162}
163
164fn benchmark_result_reducer(
165    executable: &Path,
166    config_home: &Path,
167    archive_ref: &str,
168) -> Result<Vec<Duration>, String> {
169    let text = "test routine ... ok\n".repeat(800);
170    let body = serde_json::to_vec(&json!({
171        "schemaVersion": 1,
172        "command": "cargo test",
173        "text": text,
174        "isError": false,
175        "exitCode": 0,
176        "archiveRef": archive_ref,
177        "sourceCompleteness": "unknown",
178        "preferArchiveSource": false
179    }))
180    .map_err(|error| error.to_string())?;
181    let mut request = u64::try_from(body.len())
182        .map_err(|_| "request is too large".to_owned())?
183        .to_be_bytes()
184        .to_vec();
185    request.extend_from_slice(&body);
186    let mut samples = Vec::with_capacity(SAMPLES);
187    for iteration in 0..WARMUP + SAMPLES {
188        let started = Instant::now();
189        let mut child = Command::new(executable)
190            .arg("result-reduce")
191            .env("XDG_CONFIG_HOME", config_home)
192            .stdin(Stdio::piped())
193            .stdout(Stdio::null())
194            .stderr(Stdio::piped())
195            .spawn()
196            .map_err(|error| error.to_string())?;
197        child
198            .stdin
199            .take()
200            .ok_or_else(|| "missing result reducer stdin".to_owned())?
201            .write_all(&request)
202            .map_err(|error| error.to_string())?;
203        let output = child
204            .wait_with_output()
205            .map_err(|error| error.to_string())?;
206        let elapsed = started.elapsed();
207        if !output.status.success() {
208            return Err(format!(
209                "result reducer failed: {}",
210                String::from_utf8_lossy(&output.stderr)
211            ));
212        }
213        if iteration >= WARMUP {
214            samples.push(elapsed);
215        }
216    }
217    Ok(samples)
218}
219
220fn print_stats(name: &str, samples: &[Duration]) {
221    println!("{name}_median_us: {}", percentile(samples, 50).as_micros());
222    println!("{name}_p95_us: {}", percentile(samples, 95).as_micros());
223    println!("{name}_max_us: {}", percentile(samples, 100).as_micros());
224}
225
226fn percentile(samples: &[Duration], percentile: usize) -> Duration {
227    let mut sorted = samples.to_vec();
228    sorted.sort_unstable();
229    let index = sorted
230        .len()
231        .saturating_mul(percentile)
232        .div_ceil(100)
233        .saturating_sub(1)
234        .min(sorted.len().saturating_sub(1));
235    sorted[index]
236}
indexed_output_benchmark/indexed_output_benchmark.rs

indexed_output_benchmark/
indexed_output_benchmark.rs
