1 files changed, 116 insertions, 0 deletions
diff --git a/third_party/rust/suggest/src/benchmarks/ingest.rs b/third_party/rust/suggest/src/benchmarks/ingest.rs
new file mode 100644
index 0000000000..bbefc6a00a
--- /dev/null
+++ b/third_party/rust/suggest/src/benchmarks/ingest.rs
@@ -0,0 +1,116 @@
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+use crate::{
+    benchmarks::{
+        client::{RemoteSettingsBenchmarkClient, RemoteSettingsWarmUpClient},
+        BenchmarkWithInput,
+    },
+    rs::SuggestRecordType,
+    store::SuggestStoreInner,
+    SuggestIngestionConstraints,
+};
+use std::sync::atomic::{AtomicU32, Ordering};
+
+static DB_FILE_COUNTER: AtomicU32 = AtomicU32::new(0);
+
+pub struct IngestBenchmark {
+    temp_dir: tempfile::TempDir,
+    client: RemoteSettingsBenchmarkClient,
+    record_type: SuggestRecordType,
+}
+
+impl IngestBenchmark {
+    pub fn new(record_type: SuggestRecordType) -> Self {
+        let temp_dir = tempfile::tempdir().unwrap();
+        let store = SuggestStoreInner::new(
+            temp_dir.path().join("warmup.sqlite"),
+            RemoteSettingsWarmUpClient::new(),
+        );
+        store.benchmark_ingest_records_by_type(record_type);
+        Self {
+            client: RemoteSettingsBenchmarkClient::from(store.into_settings_client()),
+            temp_dir,
+            record_type,
+        }
+    }
+}
+
+// The input for each benchmark is `SuggestStoreInner` with a fresh database.
+//
+// This is wrapped in a newtype so that it can be exposed in the public trait
+pub struct InputType(SuggestStoreInner<RemoteSettingsBenchmarkClient>);
+
+impl BenchmarkWithInput for IngestBenchmark {
+    type Input = InputType;
+
+    fn generate_input(&self) -> Self::Input {
+        let data_path = self.temp_dir.path().join(format!(
+            "db{}.sqlite",
+            DB_FILE_COUNTER.fetch_add(1, Ordering::Relaxed)
+        ));
+        let store = SuggestStoreInner::new(data_path, self.client.clone());
+        store.ensure_db_initialized();
+        InputType(store)
+    }
+
+    fn benchmarked_code(&self, input: Self::Input) {
+        let InputType(store) = input;
+        store.benchmark_ingest_records_by_type(self.record_type);
+    }
+}
+
+/// Get IngestBenchmark instances for all record types
+pub fn all_benchmarks() -> Vec<(&'static str, IngestBenchmark)> {
+    vec![
+        ("icon", IngestBenchmark::new(SuggestRecordType::Icon)),
+        (
+            "amp-wikipedia",
+            IngestBenchmark::new(SuggestRecordType::AmpWikipedia),
+        ),
+        ("amo", IngestBenchmark::new(SuggestRecordType::Amo)),
+        ("pocket", IngestBenchmark::new(SuggestRecordType::Pocket)),
+        ("yelp", IngestBenchmark::new(SuggestRecordType::Yelp)),
+        ("mdn", IngestBenchmark::new(SuggestRecordType::Mdn)),
+        ("weather", IngestBenchmark::new(SuggestRecordType::Weather)),
+        (
+            "global-config",
+            IngestBenchmark::new(SuggestRecordType::GlobalConfig),
+        ),
+        (
+            "amp-mobile",
+            IngestBenchmark::new(SuggestRecordType::AmpMobile),
+        ),
+    ]
+}
+
+pub fn print_debug_ingestion_sizes() {
+    viaduct_reqwest::use_reqwest_backend();
+    let store = SuggestStoreInner::new(
+        "file:debug_ingestion_sizes?mode=memory&cache=shared",
+        RemoteSettingsWarmUpClient::new(),
+    );
+    store
+        .ingest(SuggestIngestionConstraints::default())
+        .unwrap();
+    let table_row_counts = store.table_row_counts();
+    let client = store.into_settings_client();
+    let total_attachment_size: usize = client
+        .get_attachment_responses
+        .lock()
+        .values()
+        .map(|data| data.len())
+        .sum();
+
+    println!(
+        "Total attachment size: {}kb",
+        (total_attachment_size + 500) / 1000
+    );
+    println!();
+    println!("Database table row counts");
+    println!("-------------------------");
+    for (name, count) in table_row_counts {
+        println!("{name:30}: {count}");
+    }
+}