1 files changed, 97 insertions, 0 deletions
diff --git a/third_party/rust/suggest/src/benchmarks/client.rs b/third_party/rust/suggest/src/benchmarks/client.rs
new file mode 100644
index 0000000000..f5a21fd9cc
--- /dev/null
+++ b/third_party/rust/suggest/src/benchmarks/client.rs
@@ -0,0 +1,97 @@
+/* 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::{rs::SuggestRemoteSettingsClient, Result};
+use parking_lot::Mutex;
+use remote_settings::{Client, GetItemsOptions, RemoteSettingsConfig, RemoteSettingsResponse};
+use std::collections::HashMap;
+
+/// Remotes settings client that runs during the benchmark warm-up phase.
+///
+/// This should be used to run a full ingestion.
+/// Then it can be converted into a [RemoteSettingsBenchmarkClient], which allows benchmark code to exclude the network request time.
+/// [RemoteSettingsBenchmarkClient] implements [SuggestRemoteSettingsClient] by getting data from a HashMap rather than hitting the network.
+pub struct RemoteSettingsWarmUpClient {
+    client: Client,
+    pub get_records_responses: Mutex<HashMap<GetItemsOptions, RemoteSettingsResponse>>,
+    pub get_attachment_responses: Mutex<HashMap<String, Vec<u8>>>,
+}
+
+impl RemoteSettingsWarmUpClient {
+    pub fn new() -> Self {
+        Self {
+            client: Client::new(RemoteSettingsConfig {
+                server_url: None,
+                bucket_name: None,
+                collection_name: crate::rs::REMOTE_SETTINGS_COLLECTION.into(),
+            })
+            .unwrap(),
+            get_records_responses: Mutex::new(HashMap::new()),
+            get_attachment_responses: Mutex::new(HashMap::new()),
+        }
+    }
+}
+
+impl Default for RemoteSettingsWarmUpClient {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl SuggestRemoteSettingsClient for RemoteSettingsWarmUpClient {
+    fn get_records_with_options(
+        &self,
+        options: &GetItemsOptions,
+    ) -> Result<RemoteSettingsResponse> {
+        let response = self.client.get_records_with_options(options)?;
+        self.get_records_responses
+            .lock()
+            .insert(options.clone(), response.clone());
+        Ok(response)
+    }
+
+    fn get_attachment(&self, location: &str) -> Result<Vec<u8>> {
+        let response = self.client.get_attachment(location)?;
+        self.get_attachment_responses
+            .lock()
+            .insert(location.to_string(), response.clone());
+        Ok(response)
+    }
+}
+
+#[derive(Clone)]
+pub struct RemoteSettingsBenchmarkClient {
+    pub get_records_responses: HashMap<GetItemsOptions, RemoteSettingsResponse>,
+    pub get_attachment_responses: HashMap<String, Vec<u8>>,
+}
+
+impl SuggestRemoteSettingsClient for RemoteSettingsBenchmarkClient {
+    fn get_records_with_options(
+        &self,
+        options: &GetItemsOptions,
+    ) -> Result<RemoteSettingsResponse> {
+        Ok(self
+            .get_records_responses
+            .get(options)
+            .unwrap_or_else(|| panic!("options not found: {options:?}"))
+            .clone())
+    }
+
+    fn get_attachment(&self, location: &str) -> Result<Vec<u8>> {
+        Ok(self
+            .get_attachment_responses
+            .get(location)
+            .unwrap_or_else(|| panic!("location not found: {location:?}"))
+            .clone())
+    }
+}
+
+impl From<RemoteSettingsWarmUpClient> for RemoteSettingsBenchmarkClient {
+    fn from(warm_up_client: RemoteSettingsWarmUpClient) -> Self {
+        Self {
+            get_records_responses: warm_up_client.get_records_responses.into_inner(),
+            get_attachment_responses: warm_up_client.get_attachment_responses.into_inner(),
+        }
+    }
+}