suggest/benchmarks/
mod.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
4
5//! Benchmarking support
6//!
7//! Benchmarks are split up into two parts: the functions to be benchmarked live here, which the benchmarking code itself lives in `benches/bench.rs`.
8//! It's easier to write benchmarking code inside the main crate, where we have access to private items.
9//! However, it's easier to integrate with Cargo and criterion if benchmarks live in a separate crate.
10//!
11//! All benchmarks are defined as structs that implement either the [Benchmark] or [BenchmarkWithInput]
12
13use std::{
14    path::PathBuf,
15    sync::{
16        atomic::{AtomicU32, Ordering},
17        Mutex,
18    },
19};
20use tempfile::TempDir;
21
22use crate::{SuggestIngestionConstraints, SuggestStore};
23use remote_settings::{RemoteSettingsConfig2, RemoteSettingsContext, RemoteSettingsService};
24
25use std::sync::Arc;
26
27pub mod client;
28pub mod geoname;
29pub mod ingest;
30pub mod query;
31
32/// Trait for simple benchmarks
33///
34/// This supports simple benchmarks that don't require any input.  Note: global setup can be done
35/// in the `new()` method for the struct.
36pub trait Benchmark {
37    /// Perform the operations that we're benchmarking.
38    fn benchmarked_code(&self);
39}
40
41/// Trait for benchmarks that require input
42///
43/// This will run using Criterion's `iter_batched` function.  Criterion will create a batch of
44/// inputs, then pass each one to the benchmark's iterations.
45///
46/// This supports simple benchmarks that don't require any input.
47pub trait BenchmarkWithInput {
48    /// Input that will be created once and then passed by reference to each
49    /// of the benchmark's iterations.
50    type GlobalInput;
51
52    /// Input that will be created for each of the benchmark's iterations.
53    type IterationInput;
54
55    /// Generate the global input (not included in the benchmark time)
56    fn global_input(&self) -> Self::GlobalInput;
57
58    /// Generate the per-iteration input (not included in the benchmark time)
59    fn iteration_input(&self) -> Self::IterationInput;
60
61    /// Perform the operations that we're benchmarking.
62    fn benchmarked_code(&self, g_input: &Self::GlobalInput, i_input: Self::IterationInput);
63}
64
65fn unique_db_filename() -> String {
66    static COUNTER: AtomicU32 = AtomicU32::new(0);
67    format!("db{}.sqlite", COUNTER.fetch_add(1, Ordering::Relaxed))
68}
69
70fn unique_remote_settings_dir() -> String {
71    static COUNTER: AtomicU32 = AtomicU32::new(0);
72    format!(
73        "remote-settings-{}",
74        COUNTER.fetch_add(1, Ordering::Relaxed)
75    )
76}
77
78// Create a "starter" store that will do an initial ingest, and then
79// initialize every returned store with a copy of its DB so that each one
80// doesn't need to reingest.
81static STARTER: Mutex<Option<(TempDir, PathBuf)>> = Mutex::new(None);
82
83/// Creates a new store that will contain all provider data currently in remote
84/// settings.
85fn new_store() -> SuggestStore {
86    let mut starter = STARTER.lock().unwrap();
87    let (starter_dir, starter_db_path) = starter.get_or_insert_with(|| {
88        let temp_dir = tempfile::tempdir().unwrap();
89        let db_path = temp_dir.path().join(unique_db_filename());
90        let remote_settings_dir = temp_dir.path().join(unique_remote_settings_dir());
91        let rs_config = RemoteSettingsConfig2 {
92            bucket_name: None,
93            server: None,
94            app_context: Some(RemoteSettingsContext::default()),
95        };
96        let remote_settings_service = Arc::new(RemoteSettingsService::new(
97            remote_settings_dir.to_string_lossy().to_string(),
98            rs_config,
99        ));
100        let store = SuggestStore::new(&db_path.to_string_lossy(), remote_settings_service);
101        store
102            .ingest(SuggestIngestionConstraints::all_providers())
103            .expect("Error during ingestion");
104        store.checkpoint();
105        (temp_dir, db_path)
106    });
107
108    let db_path = starter_dir.path().join(unique_db_filename());
109    let rs_config = RemoteSettingsConfig2 {
110        bucket_name: None,
111        server: None,
112        app_context: Some(RemoteSettingsContext::default()),
113    };
114    let remote_settings_service = Arc::new(RemoteSettingsService::new("".to_string(), rs_config));
115    std::fs::copy(starter_db_path, &db_path).expect("Error copying starter DB file");
116    SuggestStore::new(&db_path.to_string_lossy(), remote_settings_service)
117}
118
119/// Cleanup the temp directory created for SuggestStore instances used in the benchmarks.
120pub fn cleanup() {
121    *STARTER.lock().unwrap() = None;
122}
suggest/benchmarks/mod.rs

suggest/benchmarks/
mod.rs
