Revision control
Copy as Markdown
Other Tools
/* 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
use std::collections::HashMap;
use crate::{db::SuggestDao, error::Error, rs, Result};
/// Remotes settings client for benchmarking
///
/// This fetches all data in `new`, then implements [rs::Client] by returning the local data.
/// Construct this one before the benchmark is run, then clone it as input for the benchmark. This
/// ensures that network time does not count towards the benchmark time.
#[derive(Clone, Default)]
pub struct RemoteSettingsBenchmarkClient {
records: Vec<rs::Record>,
attachments: HashMap<String, Vec<u8>>,
}
impl RemoteSettingsBenchmarkClient {
pub fn new() -> Result<Self> {
let mut new_benchmark_client = Self::default();
new_benchmark_client.fetch_data_with_client(
remote_settings::RemoteSettings::new(remote_settings::RemoteSettingsConfig {
server: None,
bucket_name: None,
collection_name: "quicksuggest".to_owned(),
server_url: None,
})?,
rs::Collection::Quicksuggest,
)?;
new_benchmark_client.fetch_data_with_client(
remote_settings::RemoteSettings::new(remote_settings::RemoteSettingsConfig {
server: None,
bucket_name: None,
collection_name: "fakespot-suggest-products".to_owned(),
server_url: None,
})?,
rs::Collection::Fakespot,
)?;
Ok(new_benchmark_client)
}
fn fetch_data_with_client(
&mut self,
client: remote_settings::RemoteSettings,
collection: rs::Collection,
) -> Result<()> {
let response = client.get_records()?;
for r in &response.records {
if let Some(a) = &r.attachment {
self.attachments
.insert(a.location.clone(), client.get_attachment(&a.location)?);
}
}
self.records.extend(
response
.records
.into_iter()
.filter_map(|r| rs::Record::new(r, collection).ok()),
);
Ok(())
}
pub fn attachment_size_by_record_type(&self) -> Vec<(rs::SuggestRecordType, usize)> {
let mut sizes = HashMap::<rs::SuggestRecordType, usize>::new();
for record in self.records.iter() {
let record_type = rs::SuggestRecordType::from(&record.payload);
if let Some(a) = &record.attachment {
if let Some(attachment) = self.attachments.get(&a.location) {
sizes
.entry(record_type)
.and_modify(|size| *size += attachment.len())
.or_insert(attachment.len());
}
}
}
let mut sizes_vec: Vec<_> = sizes.into_iter().collect();
sizes_vec.sort_by_key(|(_, size)| *size);
sizes_vec.reverse();
sizes_vec
}
pub fn total_attachment_size(&self) -> usize {
self.attachments.values().map(|a| a.len()).sum()
}
}
impl rs::Client for RemoteSettingsBenchmarkClient {
fn get_records(
&self,
collection: rs::Collection,
_db: &mut SuggestDao,
) -> Result<Vec<rs::Record>> {
Ok(self
.records
.iter()
.filter(|r| r.collection == collection)
.cloned()
.collect())
}
fn download_attachment(&self, record: &rs::Record) -> Result<Vec<u8>> {
match &record.attachment {
Some(a) => match self.attachments.get(&a.location) {
Some(data) => Ok(data.clone()),
None => Err(Error::MissingAttachment(record.id.to_string())),
},
None => Err(Error::MissingAttachment(record.id.to_string())),
}
}
}