Skip to main content

open_geocode/
bench.rs

1use std::{
2    fs,
3    path::{Path, PathBuf},
4    sync::Arc,
5    time::Instant,
6};
7
8use anyhow::{Context, Result};
9use serde::{Deserialize, Serialize};
10
11use crate::{
12    builder::report::{BuilderReport, PhaseTimings, ScratchReport, Throughput},
13    pack::{AUDIT_DIR, BUILD_REPORT_FILE, PackManifest, PackReader},
14    reverse::{PackReverseGeocoder, ReverseGeocodeOptions},
15    search::{PackTextSearcher, TextAutocompleteOptions, TextSearchOptions},
16};
17
18#[derive(Debug, Clone)]
19pub struct PackBenchmarkOptions {
20    pub pack: PathBuf,
21    pub queries: Option<PathBuf>,
22    pub iterations: usize,
23    pub warmup: usize,
24}
25
26#[derive(Debug, Serialize)]
27pub struct PackBenchmarkReport {
28    pub settings: PackBenchmarkSettings,
29    pub pack: PackMetricReport,
30    pub open: OpenBenchmarkReport,
31    pub queries: QueryBenchmarkReport,
32}
33
34#[derive(Debug, Serialize)]
35pub struct PackBenchmarkSettings {
36    pub pack: String,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub queries: Option<String>,
39    pub iterations: usize,
40    pub warmup: usize,
41}
42
43#[derive(Debug, Serialize)]
44pub struct PackMetricReport {
45    pub manifest: PackManifest,
46    pub bytes: PackByteMetrics,
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub build: Option<BuildMetricReport>,
49}
50
51/// Build timings and throughput from the build report next to the Pack.
52#[derive(Debug, Serialize)]
53pub struct BuildMetricReport {
54    pub input_bytes: u64,
55    pub accepted_records: u64,
56    pub rejected_records: u64,
57    pub total_seconds: f64,
58    #[serde(skip_serializing_if = "Option::is_none")]
59    pub input_mib_per_sec: Option<f64>,
60    pub phases: PhaseTimings,
61    pub throughput: Throughput,
62    pub scratch: ScratchReport,
63}
64
65/// Pack bytes by section group.
66#[derive(Debug, Default, Serialize)]
67pub struct PackByteMetrics {
68    pub total: u64,
69    pub records: u64,
70    pub context: u64,
71    pub text_index: u64,
72    pub spatial_index: u64,
73    pub other: u64,
74    #[serde(skip_serializing_if = "Option::is_none")]
75    pub bytes_per_record: Option<f64>,
76    #[serde(skip_serializing_if = "Option::is_none")]
77    pub records_bytes_per_record: Option<f64>,
78    #[serde(skip_serializing_if = "Option::is_none")]
79    pub text_index_bytes_per_record: Option<f64>,
80    #[serde(skip_serializing_if = "Option::is_none")]
81    pub spatial_index_bytes_per_record: Option<f64>,
82}
83
84#[derive(Debug, Serialize)]
85pub struct OpenBenchmarkReport {
86    pub pack_reader_ms: f64,
87    pub text_searcher_ms: f64,
88    pub reverse_geocoder_ms: f64,
89}
90
91#[derive(Debug, Serialize)]
92pub struct QueryBenchmarkReport {
93    pub search: OperationBenchmarkReport<TextQueryCaseReport>,
94    pub autocomplete: OperationBenchmarkReport<TextQueryCaseReport>,
95    pub reverse: OperationBenchmarkReport<ReverseQueryCaseReport>,
96}
97
98#[derive(Debug, Serialize)]
99pub struct OperationBenchmarkReport<T> {
100    pub case_count: usize,
101    pub measured_runs: usize,
102    pub warmup_runs: usize,
103    #[serde(skip_serializing_if = "Option::is_none")]
104    pub latency: Option<LatencyStats>,
105    pub cases: Vec<T>,
106}
107
108#[derive(Debug, Serialize)]
109pub struct TextQueryCaseReport {
110    #[serde(skip_serializing_if = "Option::is_none")]
111    pub name: Option<String>,
112    pub query: String,
113    pub limit: usize,
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub layer: Option<String>,
116    pub hit_count: usize,
117    pub latency: LatencyStats,
118}
119
120#[derive(Debug, Serialize)]
121pub struct ReverseQueryCaseReport {
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub name: Option<String>,
124    pub lon: f64,
125    pub lat: f64,
126    pub result_present: bool,
127    pub latency: LatencyStats,
128}
129
130#[derive(Debug, Clone, Serialize, PartialEq)]
131pub struct LatencyStats {
132    pub min_ms: f64,
133    pub p50_ms: f64,
134    pub p90_ms: f64,
135    pub p95_ms: f64,
136    pub p99_ms: f64,
137    pub max_ms: f64,
138    pub mean_ms: f64,
139    pub total_ms: f64,
140}
141
142#[derive(Debug, Default, Deserialize)]
143struct BenchmarkFixture {
144    #[serde(default)]
145    search: Vec<TextQueryFixture>,
146    #[serde(default)]
147    autocomplete: Vec<TextQueryFixture>,
148    #[serde(default)]
149    reverse: Vec<ReverseQueryFixture>,
150}
151
152#[derive(Debug, Clone, Deserialize)]
153struct TextQueryFixture {
154    #[serde(default)]
155    name: Option<String>,
156    #[serde(alias = "q")]
157    query: String,
158    #[serde(default)]
159    limit: Option<usize>,
160    #[serde(default)]
161    layer: Option<String>,
162}
163
164#[derive(Debug, Clone, Deserialize)]
165struct ReverseQueryFixture {
166    #[serde(default)]
167    name: Option<String>,
168    lon: f64,
169    lat: f64,
170}
171
172pub fn benchmark_pack(options: PackBenchmarkOptions) -> Result<PackBenchmarkReport> {
173    let iterations = options.iterations.max(1);
174    let warmup = options.warmup;
175    let fixture = read_fixture(options.queries.as_deref())?;
176
177    let (reader, pack_reader_ms) = measure_value(|| PackReader::open(&options.pack).map(Arc::new))?;
178    let pack = pack_metrics(&reader)?;
179    let (searcher, text_searcher_ms) =
180        measure_value(|| PackTextSearcher::from_pack(Arc::clone(&reader)))?;
181    let (reverse_geocoder, reverse_geocoder_ms) =
182        measure_value(|| PackReverseGeocoder::from_pack(Arc::clone(&reader)))?;
183
184    let queries = QueryBenchmarkReport {
185        search: benchmark_search_cases(&searcher, &fixture.search, iterations, warmup)?,
186        autocomplete: benchmark_autocomplete_cases(
187            &searcher,
188            &fixture.autocomplete,
189            iterations,
190            warmup,
191        )?,
192        reverse: benchmark_reverse_cases(&reverse_geocoder, &fixture.reverse, iterations, warmup)?,
193    };
194
195    Ok(PackBenchmarkReport {
196        settings: PackBenchmarkSettings {
197            pack: options.pack.display().to_string(),
198            queries: options
199                .queries
200                .as_ref()
201                .map(|path| path.display().to_string()),
202            iterations,
203            warmup,
204        },
205        pack,
206        open: OpenBenchmarkReport {
207            pack_reader_ms,
208            text_searcher_ms,
209            reverse_geocoder_ms,
210        },
211        queries,
212    })
213}
214
215fn read_fixture(path: Option<&Path>) -> Result<BenchmarkFixture> {
216    let Some(path) = path else {
217        return Ok(BenchmarkFixture::default());
218    };
219    let file =
220        fs::File::open(path).with_context(|| format!("failed to open {}", path.display()))?;
221    serde_json::from_reader(file).with_context(|| format!("failed to parse {}", path.display()))
222}
223
224fn pack_metrics(reader: &PackReader) -> Result<PackMetricReport> {
225    let manifest = reader.manifest();
226    let mut bytes = PackByteMetrics {
227        total: reader.container().file_size(),
228        ..PackByteMetrics::default()
229    };
230    for (name, len) in reader.section_sizes() {
231        let group = match name.split('/').next().unwrap_or_default() {
232            "records" => &mut bytes.records,
233            "context" => &mut bytes.context,
234            "text" => &mut bytes.text_index,
235            "spatial" => &mut bytes.spatial_index,
236            _ => &mut bytes.other,
237        };
238        *group += len;
239    }
240    if manifest.record_count > 0 {
241        let count = manifest.record_count as f64;
242        bytes.bytes_per_record = Some(bytes.total as f64 / count);
243        bytes.records_bytes_per_record = Some(bytes.records as f64 / count);
244        bytes.text_index_bytes_per_record = Some(bytes.text_index as f64 / count);
245        bytes.spatial_index_bytes_per_record = Some(bytes.spatial_index as f64 / count);
246    }
247    Ok(PackMetricReport {
248        manifest: manifest.clone(),
249        bytes,
250        build: read_build_report(reader.path())?.map(build_metrics),
251    })
252}
253
254fn read_build_report(pack_file: &Path) -> Result<Option<BuilderReport>> {
255    let Some(path) = pack_file
256        .parent()
257        .map(|dir| dir.join(AUDIT_DIR).join(BUILD_REPORT_FILE))
258        .filter(|path| path.is_file())
259    else {
260        return Ok(None);
261    };
262    let file =
263        fs::File::open(&path).with_context(|| format!("failed to open {}", path.display()))?;
264    serde_json::from_reader(file)
265        .map(Some)
266        .with_context(|| format!("failed to parse {}", path.display()))
267}
268
269fn build_metrics(report: BuilderReport) -> BuildMetricReport {
270    let total_seconds = report.phases.total_ms as f64 / 1_000.0;
271    BuildMetricReport {
272        input_bytes: report.input_bytes,
273        accepted_records: report.accepted.total,
274        rejected_records: report.rejected.total,
275        total_seconds,
276        input_mib_per_sec: (total_seconds > 0.0)
277            .then(|| report.input_bytes as f64 / 1_048_576.0 / total_seconds),
278        phases: report.phases,
279        throughput: report.throughput,
280        scratch: report.scratch,
281    }
282}
283
284fn benchmark_search_cases(
285    searcher: &PackTextSearcher,
286    cases: &[TextQueryFixture],
287    iterations: usize,
288    warmup: usize,
289) -> Result<OperationBenchmarkReport<TextQueryCaseReport>> {
290    let mut reports = Vec::new();
291    let mut all_durations = Vec::new();
292    for case in cases {
293        let limit = case.limit.unwrap_or(10);
294        let mut hit_count = 0;
295        let durations = measure_iterations(iterations, warmup, || {
296            let hits = searcher.search(TextSearchOptions {
297                query: case.query.clone(),
298                limit,
299                layer: case.layer.clone(),
300            })?;
301            hit_count = hits.len();
302            Ok(())
303        })?;
304        all_durations.extend(durations.iter().copied());
305        reports.push(TextQueryCaseReport {
306            name: case.name.clone(),
307            query: case.query.clone(),
308            limit,
309            layer: case.layer.clone(),
310            hit_count,
311            latency: LatencyStats::from_nanos(&durations),
312        });
313    }
314
315    Ok(operation_report(
316        cases.len(),
317        iterations,
318        warmup,
319        reports,
320        &all_durations,
321    ))
322}
323
324fn benchmark_autocomplete_cases(
325    searcher: &PackTextSearcher,
326    cases: &[TextQueryFixture],
327    iterations: usize,
328    warmup: usize,
329) -> Result<OperationBenchmarkReport<TextQueryCaseReport>> {
330    let mut reports = Vec::new();
331    let mut all_durations = Vec::new();
332    for case in cases {
333        let limit = case.limit.unwrap_or(10);
334        let mut hit_count = 0;
335        let durations = measure_iterations(iterations, warmup, || {
336            let hits = searcher.autocomplete(TextAutocompleteOptions {
337                query: case.query.clone(),
338                limit,
339                layer: case.layer.clone(),
340            })?;
341            hit_count = hits.len();
342            Ok(())
343        })?;
344        all_durations.extend(durations.iter().copied());
345        reports.push(TextQueryCaseReport {
346            name: case.name.clone(),
347            query: case.query.clone(),
348            limit,
349            layer: case.layer.clone(),
350            hit_count,
351            latency: LatencyStats::from_nanos(&durations),
352        });
353    }
354
355    Ok(operation_report(
356        cases.len(),
357        iterations,
358        warmup,
359        reports,
360        &all_durations,
361    ))
362}
363
364fn benchmark_reverse_cases(
365    geocoder: &PackReverseGeocoder,
366    cases: &[ReverseQueryFixture],
367    iterations: usize,
368    warmup: usize,
369) -> Result<OperationBenchmarkReport<ReverseQueryCaseReport>> {
370    let mut reports = Vec::new();
371    let mut all_durations = Vec::new();
372    for case in cases {
373        let mut result_present = false;
374        let durations = measure_iterations(iterations, warmup, || {
375            let response = geocoder.reverse(ReverseGeocodeOptions {
376                lon: case.lon,
377                lat: case.lat,
378            })?;
379            result_present = response.result.is_some();
380            Ok(())
381        })?;
382        all_durations.extend(durations.iter().copied());
383        reports.push(ReverseQueryCaseReport {
384            name: case.name.clone(),
385            lon: case.lon,
386            lat: case.lat,
387            result_present,
388            latency: LatencyStats::from_nanos(&durations),
389        });
390    }
391
392    Ok(operation_report(
393        cases.len(),
394        iterations,
395        warmup,
396        reports,
397        &all_durations,
398    ))
399}
400
401fn operation_report<T>(
402    case_count: usize,
403    iterations: usize,
404    warmup: usize,
405    cases: Vec<T>,
406    durations: &[u128],
407) -> OperationBenchmarkReport<T> {
408    OperationBenchmarkReport {
409        case_count,
410        measured_runs: case_count * iterations,
411        warmup_runs: case_count * warmup,
412        latency: (!durations.is_empty()).then(|| LatencyStats::from_nanos(durations)),
413        cases,
414    }
415}
416
417fn measure_value<T>(measure: impl FnOnce() -> Result<T>) -> Result<(T, f64)> {
418    let started = Instant::now();
419    let value = measure()?;
420    Ok((value, nanos_to_ms(started.elapsed().as_nanos())))
421}
422
423fn measure_iterations(
424    iterations: usize,
425    warmup: usize,
426    mut measure: impl FnMut() -> Result<()>,
427) -> Result<Vec<u128>> {
428    for _ in 0..warmup {
429        measure()?;
430    }
431
432    let mut durations = Vec::with_capacity(iterations);
433    for _ in 0..iterations {
434        let started = Instant::now();
435        measure()?;
436        durations.push(started.elapsed().as_nanos());
437    }
438    Ok(durations)
439}
440
441impl LatencyStats {
442    fn from_nanos(durations: &[u128]) -> Self {
443        debug_assert!(!durations.is_empty());
444        let mut sorted = durations.to_vec();
445        sorted.sort_unstable();
446        let total = sorted.iter().sum::<u128>();
447        let mean = total as f64 / sorted.len() as f64;
448        Self {
449            min_ms: nanos_to_ms(*sorted.first().expect("duration")),
450            p50_ms: nanos_to_ms(percentile(&sorted, 50.0)),
451            p90_ms: nanos_to_ms(percentile(&sorted, 90.0)),
452            p95_ms: nanos_to_ms(percentile(&sorted, 95.0)),
453            p99_ms: nanos_to_ms(percentile(&sorted, 99.0)),
454            max_ms: nanos_to_ms(*sorted.last().expect("duration")),
455            mean_ms: mean / 1_000_000.0,
456            total_ms: nanos_to_ms(total),
457        }
458    }
459}
460
461fn percentile(sorted: &[u128], percentile: f64) -> u128 {
462    let rank = ((percentile / 100.0) * sorted.len() as f64).ceil() as usize;
463    let index = rank.saturating_sub(1).min(sorted.len() - 1);
464    sorted[index]
465}
466
467fn nanos_to_ms(nanos: u128) -> f64 {
468    nanos as f64 / 1_000_000.0
469}
470
471#[cfg(test)]
472mod tests {
473    use crate::{
474        pack::PackWriter,
475        record::{
476            AddressComponents, AddressRecord, LocationPrecision, OsmObjectType, SourceProvenance,
477            point_geometry,
478        },
479    };
480
481    use super::*;
482
483    #[test]
484    fn reports_pack_metrics_without_query_fixture() {
485        let temp_dir = temp_pack_path("bench-metrics");
486        let _ = fs::remove_dir_all(&temp_dir);
487        write_test_pack(&temp_dir);
488
489        let report = benchmark_pack(PackBenchmarkOptions {
490            pack: temp_dir.clone(),
491            queries: None,
492            iterations: 2,
493            warmup: 1,
494        })
495        .expect("benchmark");
496
497        assert_eq!(report.pack.manifest.record_count, 1);
498        assert!(report.pack.bytes.total > 0);
499        assert!(report.pack.bytes.records > 0);
500        assert!(report.pack.bytes.text_index > 0);
501        assert!(report.pack.bytes.spatial_index > 0);
502        assert!(
503            report.pack.build.is_none(),
504            "a Pack written directly has no build report"
505        );
506        assert!(report.open.pack_reader_ms >= 0.0);
507        assert_eq!(report.queries.search.case_count, 0);
508
509        let _ = fs::remove_dir_all(temp_dir);
510    }
511
512    #[test]
513    fn benchmarks_query_fixture_cases() {
514        let temp_dir = temp_pack_path("bench-queries");
515        let _ = fs::remove_dir_all(&temp_dir);
516        write_test_pack(&temp_dir);
517
518        let fixture_path = temp_dir.join("queries.json");
519        fs::write(
520            &fixture_path,
521            r#"{
522              "search": [{"name": "king", "q": "King Street Toronto", "limit": 5}],
523              "autocomplete": [{"name": "prefix", "q": "kin", "limit": 5}],
524              "reverse": [{"name": "point", "lon": -79.4, "lat": 43.6}]
525            }"#,
526        )
527        .expect("write fixture");
528
529        let report = benchmark_pack(PackBenchmarkOptions {
530            pack: temp_dir.clone(),
531            queries: Some(fixture_path),
532            iterations: 2,
533            warmup: 1,
534        })
535        .expect("benchmark");
536
537        assert_eq!(report.queries.search.case_count, 1);
538        assert_eq!(report.queries.search.measured_runs, 2);
539        assert_eq!(report.queries.search.warmup_runs, 1);
540        assert_eq!(report.queries.search.cases[0].hit_count, 1);
541        assert_eq!(report.queries.autocomplete.cases[0].hit_count, 1);
542        assert!(report.queries.reverse.cases[0].result_present);
543        assert!(report.queries.search.latency.is_some());
544
545        let _ = fs::remove_dir_all(temp_dir);
546    }
547
548    fn write_test_pack(path: &Path) {
549        let mut writer = PackWriter::create(path).expect("writer");
550        writer
551            .write(
552                &AddressRecord {
553                    address: AddressComponents {
554                        number: "10".to_string(),
555                        street: Some("King Street".to_string()),
556                        place: None,
557                        unit: None,
558                        locality: Some("Toronto".to_string()),
559                        region: Some("Ontario".to_string()),
560                        postcode: Some("M5V 1A1".to_string()),
561                        country: Some("CA".to_string()),
562                    },
563                    geometry: point_geometry(-79.4, 43.6),
564                    location_precision: LocationPrecision::Point,
565                    source: SourceProvenance::osm(OsmObjectType::Node, 1),
566                }
567                .into(),
568                None,
569            )
570            .expect("write address");
571        writer.finish().expect("finish");
572    }
573
574    fn temp_pack_path(name: &str) -> PathBuf {
575        std::env::temp_dir().join(format!("open-geocode-{name}-{}", std::process::id()))
576    }
577}