suggest/benchmarks/
client.rs

Help
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        new_benchmark_client.fetch_data_with_client(
42            remote_settings::RemoteSettings::new(remote_settings::RemoteSettingsConfig {
43                server: None,
44                bucket_name: None,
45                collection_name: rs::Collection::Fakespot.name().to_owned(),
46                server_url: None,
47            })?,
48            rs::Collection::Fakespot,
49        )?;
50        Ok(new_benchmark_client)
51    }
52
53    fn fetch_data_with_client(
54        &mut self,
55        client: remote_settings::RemoteSettings,
56        collection: rs::Collection,
57    ) -> Result<()> {
58        let response = client.get_records()?;
59        for r in &response.records {
60            if let Some(a) = &r.attachment {
61                self.attachments
62                    .insert(a.location.clone(), client.get_attachment(&a.location)?);
63            }
64        }
65        self.records.extend(
66            response
67                .records
68                .into_iter()
69                .filter_map(|r| rs::Record::new(r, collection).ok()),
70        );
71        Ok(())
72    }
73
74    pub fn attachment_size_by_record_type(&self) -> Vec<(rs::SuggestRecordType, usize)> {
75        let mut sizes = HashMap::<rs::SuggestRecordType, usize>::new();
76        for record in self.records.iter() {
77            let record_type = rs::SuggestRecordType::from(&record.payload);
78            if let Some(a) = &record.attachment {
79                if let Some(attachment) = self.attachments.get(&a.location) {
80                    sizes
81                        .entry(record_type)
82                        .and_modify(|size| *size += attachment.len())
83                        .or_insert(attachment.len());
84                }
85            }
86        }
87        let mut sizes_vec: Vec<_> = sizes.into_iter().collect();
88        sizes_vec.sort_by_key(|(_, size)| *size);
89        sizes_vec.reverse();
90        sizes_vec
91    }
92
93    pub fn total_attachment_size(&self) -> usize {
94        self.attachments.values().map(|a| a.len()).sum()
95    }
96}
97
98impl rs::Client for RemoteSettingsBenchmarkClient {
99    fn get_records(&self, collection: rs::Collection) -> Result<Vec<rs::Record>> {
100        Ok(self
101            .records
102            .iter()
103            .filter(|r| r.collection == collection)
104            .cloned()
105            .collect())
106    }
107
108    fn download_attachment(&self, record: &rs::Record) -> Result<Vec<u8>> {
109        match &record.attachment {
110            Some(a) => match self.attachments.get(&a.location) {
111                Some(data) => Ok(data.clone()),
112                None => Err(Error::MissingAttachment(record.id.to_string())),
113            },
114            None => Err(Error::MissingAttachment(record.id.to_string())),
115        }
116    }
117}
suggest/benchmarks/client.rs

suggest/benchmarks/
client.rs
