suggest/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
5use std::collections::HashMap;
6
7use crate::{error::Error, rs, Result};
8
9/// Remotes settings client for benchmarking
10///
11/// This fetches all data in `new`, then implements [rs::Client] by returning the local data.
12/// Construct this one before the benchmark is run, then clone it as input for the benchmark.  This
13/// ensures that network time does not count towards the benchmark time.
14#[derive(Clone, Default)]
15pub struct RemoteSettingsBenchmarkClient {
16    records: Vec<rs::Record>,
17    attachments: HashMap<String, Vec<u8>>,
18}
19
20impl RemoteSettingsBenchmarkClient {
21    pub fn new() -> Result<Self> {
22        let mut new_benchmark_client = Self::default();
23        new_benchmark_client.fetch_data_with_client(
24            remote_settings::RemoteSettings::new(remote_settings::RemoteSettingsConfig {
25                server: None,
26                bucket_name: None,
27                collection_name: rs::Collection::Amp.name().to_owned(),
28                server_url: None,
29            })?,
30            rs::Collection::Amp,
31        )?;
32        new_benchmark_client.fetch_data_with_client(
33            remote_settings::RemoteSettings::new(remote_settings::RemoteSettingsConfig {
34                server: None,
35                bucket_name: None,
36                collection_name: rs::Collection::Other.name().to_owned(),
37                server_url: None,
38            })?,
39            rs::Collection::Other,
40        )?;
41        Ok(new_benchmark_client)
42    }
43
44    fn fetch_data_with_client(
45        &mut self,
46        client: remote_settings::RemoteSettings,
47        collection: rs::Collection,
48    ) -> Result<()> {
49        let response = client.get_records()?;
50        for r in &response.records {
51            if let Some(a) = &r.attachment {
52                self.attachments
53                    .insert(a.location.clone(), client.get_attachment(&a.location)?);
54            }
55        }
56        self.records.extend(
57            response
58                .records
59                .into_iter()
60                .filter_map(|r| rs::Record::new(r, collection).ok()),
61        );
62        Ok(())
63    }
64
65    pub fn attachment_size_by_record_type(&self) -> Vec<(rs::SuggestRecordType, usize)> {
66        let mut sizes = HashMap::<rs::SuggestRecordType, usize>::new();
67        for record in self.records.iter() {
68            let record_type = rs::SuggestRecordType::from(&record.payload);
69            if let Some(a) = &record.attachment {
70                if let Some(attachment) = self.attachments.get(&a.location) {
71                    sizes
72                        .entry(record_type)
73                        .and_modify(|size| *size += attachment.len())
74                        .or_insert(attachment.len());
75                }
76            }
77        }
78        let mut sizes_vec: Vec<_> = sizes.into_iter().collect();
79        sizes_vec.sort_by_key(|(_, size)| *size);
80        sizes_vec.reverse();
81        sizes_vec
82    }
83
84    pub fn total_attachment_size(&self) -> usize {
85        self.attachments.values().map(|a| a.len()).sum()
86    }
87}
88
89impl rs::Client for RemoteSettingsBenchmarkClient {
90    fn get_records(&self, collection: rs::Collection) -> Result<Vec<rs::Record>> {
91        Ok(self
92            .records
93            .iter()
94            .filter(|r| r.collection == collection)
95            .cloned()
96            .collect())
97    }
98
99    fn download_attachment(&self, record: &rs::Record) -> Result<Vec<u8>> {
100        match &record.attachment {
101            Some(a) => match self.attachments.get(&a.location) {
102                Some(data) => Ok(data.clone()),
103                None => Err(Error::MissingAttachment(record.id.to_string())),
104            },
105            None => Err(Error::MissingAttachment(record.id.to_string())),
106        }
107    }
108}
suggest/benchmarks/client.rs

suggest/benchmarks/
client.rs
