third_party/rust/suggest/src/benchmarks/client.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 use crate::{rs::SuggestRemoteSettingsClient, Result};
   6 use parking_lot::Mutex;
   7 use remote_settings::{Client, GetItemsOptions, RemoteSettingsConfig, RemoteSettingsResponse};
   8 use std::collections::HashMap;
   9
  10 /// Remotes settings client that runs during the benchmark warm-up phase.
  11 ///
  12 /// This should be used to run a full ingestion.
  13 /// Then it can be converted into a [RemoteSettingsBenchmarkClient], which allows benchmark code to exclude the network request time.
  14 /// [RemoteSettingsBenchmarkClient] implements [SuggestRemoteSettingsClient] by getting data from a HashMap rather than hitting the network.
  15 pub struct RemoteSettingsWarmUpClient {
  16     client: Client,
  17     pub get_records_responses: Mutex<HashMap<GetItemsOptions, RemoteSettingsResponse>>,
  18     pub get_attachment_responses: Mutex<HashMap<String, Vec<u8>>>,
  19 }
  20
  21 impl RemoteSettingsWarmUpClient {
  22     pub fn new() -> Self {
  23         Self {
  24             client: Client::new(RemoteSettingsConfig {
  25                 server_url: None,
  26                 bucket_name: None,
  27                 collection_name: crate::rs::REMOTE_SETTINGS_COLLECTION.into(),
  28             })
  29             .unwrap(),
  30             get_records_responses: Mutex::new(HashMap::new()),
  31             get_attachment_responses: Mutex::new(HashMap::new()),
  32         }
  33     }
  34 }
  35
  36 impl Default for RemoteSettingsWarmUpClient {
  37     fn default() -> Self {
  38         Self::new()
  39     }
  40 }
  41
  42 impl SuggestRemoteSettingsClient for RemoteSettingsWarmUpClient {
  43     fn get_records_with_options(
  44         &self,
  45         options: &GetItemsOptions,
  46     ) -> Result<RemoteSettingsResponse> {
  47         let response = self.client.get_records_with_options(options)?;
  48         self.get_records_responses
  49             .lock()
  50             .insert(options.clone(), response.clone());
  51         Ok(response)
  52     }
  53
  54     fn get_attachment(&self, location: &str) -> Result<Vec<u8>> {
  55         let response = self.client.get_attachment(location)?;
  56         self.get_attachment_responses
  57             .lock()
  58             .insert(location.to_string(), response.clone());
  59         Ok(response)
  60     }
  61 }
  62
  63 #[derive(Clone)]
  64 pub struct RemoteSettingsBenchmarkClient {
  65     pub get_records_responses: HashMap<GetItemsOptions, RemoteSettingsResponse>,
  66     pub get_attachment_responses: HashMap<String, Vec<u8>>,
  67 }
  68
  69 impl SuggestRemoteSettingsClient for RemoteSettingsBenchmarkClient {
  70     fn get_records_with_options(
  71         &self,
  72         options: &GetItemsOptions,
  73     ) -> Result<RemoteSettingsResponse> {
  74         Ok(self
  75             .get_records_responses
  76             .get(options)
  77             .unwrap_or_else(|| panic!("options not found: {options:?}"))
  78             .clone())
  79     }
  80
  81     fn get_attachment(&self, location: &str) -> Result<Vec<u8>> {
  82         Ok(self
  83             .get_attachment_responses
  84             .get(location)
  85             .unwrap_or_else(|| panic!("location not found: {location:?}"))
  86             .clone())
  87     }
  88 }
  89
  90 impl From<RemoteSettingsWarmUpClient> for RemoteSettingsBenchmarkClient {
  91     fn from(warm_up_client: RemoteSettingsWarmUpClient) -> Self {
  92         Self {
  93             get_records_responses: warm_up_client.get_records_responses.into_inner(),
  94             get_attachment_responses: warm_up_client.get_attachment_responses.into_inner(),
  95         }
  96     }
  97 }