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

suggest/benchmarks/
client.rs
