glean_core/storage/
mod.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 https://mozilla.org/MPL/2.0/.
4
5#![allow(non_upper_case_globals)]
6
7//! Storage snapshotting.
8
9use std::collections::HashMap;
10
11use serde_json::{json, Value as JsonValue};
12
13use crate::coverage::record_coverage;
14use crate::database::Database;
15use crate::metrics::dual_labeled_counter::RECORD_SEPARATOR;
16use crate::metrics::Metric;
17use crate::Lifetime;
18
19// An internal ping name, not to be touched by anything else
20pub(crate) const INTERNAL_STORAGE: &str = "glean_internal_info";
21
22/// Snapshot metrics from the underlying database.
23pub struct StorageManager;
24
25/// Labeled metrics are stored as `<metric id>/<label>`.
26/// They need to go into a nested object in the final snapshot.
27///
28/// We therefore extract the metric id and the label from the key and construct the new object or
29/// add to it.
30fn snapshot_labeled_metrics(
31    snapshot: &mut HashMap<String, HashMap<String, JsonValue>>,
32    metric_id: &str,
33    metric: &Metric,
34) {
35    // Explicit match for supported labeled metrics, avoiding the formatting string
36    let ping_section = match metric.ping_section() {
37        "boolean" => "labeled_boolean".to_string(),
38        "counter" => "labeled_counter".to_string(),
39        "timing_distribution" => "labeled_timing_distribution".to_string(),
40        "memory_distribution" => "labeled_memory_distribution".to_string(),
41        "custom_distribution" => "labeled_custom_distribution".to_string(),
42        "quantity" => "labeled_quantity".to_string(),
43        // This should never happen, we covered all cases.
44        // Should we ever extend it this would however at least catch it and do the right thing.
45        _ => format!("labeled_{}", metric.ping_section()),
46    };
47    let map = snapshot.entry(ping_section).or_default();
48
49    // Safe unwrap, the function is only called when the id does contain a '/'
50    let (metric_id, label) = metric_id.split_once('/').unwrap();
51
52    let obj = map.entry(metric_id.into()).or_insert_with(|| json!({}));
53    let obj = obj.as_object_mut().unwrap(); // safe unwrap, we constructed the object above
54    obj.insert(label.into(), metric.as_json());
55}
56
57/// Dual Labeled metrics are stored as `<metric id><\x1e><key><\x1e><category>`.
58/// They need to go into a nested object in the final snapshot.
59///
60/// We therefore extract the metric id and the label from the key and construct the new object or
61/// add to it.
62fn snapshot_dual_labeled_metrics(
63    snapshot: &mut HashMap<String, HashMap<String, JsonValue>>,
64    metric_id: &str,
65    metric: &Metric,
66) {
67    let ping_section = format!("dual_labeled_{}", metric.ping_section());
68    let map = snapshot.entry(ping_section).or_default();
69    let parts = metric_id.split(RECORD_SEPARATOR).collect::<Vec<&str>>();
70
71    let obj = map
72        .entry(parts[0].into())
73        .or_insert_with(|| json!({}))
74        .as_object_mut()
75        .unwrap(); // safe unwrap, we constructed the object above
76    let key_obj = obj.entry(parts[1].to_string()).or_insert_with(|| json!({}));
77    let key_obj = key_obj.as_object_mut().unwrap();
78    key_obj.insert(parts[2].into(), metric.as_json());
79}
80
81impl StorageManager {
82    /// Snapshots the given store and optionally clear it.
83    ///
84    /// # Arguments
85    ///
86    /// * `storage` - the database to read from.
87    /// * `store_name` - the store to snapshot.
88    /// * `clear_store` - whether to clear the data after snapshotting.
89    ///
90    /// # Returns
91    ///
92    /// The stored data in a string encoded as JSON.
93    /// If no data for the store exists, `None` is returned.
94    pub fn snapshot(
95        &self,
96        storage: &Database,
97        store_name: &str,
98        clear_store: bool,
99    ) -> Option<String> {
100        self.snapshot_as_json(storage, store_name, clear_store)
101            .map(|data| ::serde_json::to_string_pretty(&data).unwrap())
102    }
103
104    /// Snapshots the given store and optionally clear it.
105    ///
106    /// # Arguments
107    ///
108    /// * `storage` - the database to read from.
109    /// * `store_name` - the store to snapshot.
110    /// * `clear_store` - whether to clear the data after snapshotting.
111    ///
112    /// # Returns
113    ///
114    /// A JSON representation of the stored data.
115    /// If no data for the store exists, `None` is returned.
116    pub fn snapshot_as_json(
117        &self,
118        storage: &Database,
119        store_name: &str,
120        clear_store: bool,
121    ) -> Option<JsonValue> {
122        let mut snapshot: HashMap<String, HashMap<String, JsonValue>> = HashMap::new();
123
124        let mut snapshotter = |metric_id: &[u8], metric: &Metric| {
125            let metric_id = String::from_utf8_lossy(metric_id).into_owned();
126            if metric_id.contains('/') {
127                snapshot_labeled_metrics(&mut snapshot, &metric_id, metric);
128            } else if metric_id.split(RECORD_SEPARATOR).count() == 3 {
129                snapshot_dual_labeled_metrics(&mut snapshot, &metric_id, metric);
130            } else {
131                let map = snapshot.entry(metric.ping_section().into()).or_default();
132                map.insert(metric_id, metric.as_json());
133            }
134        };
135
136        storage.iter_store_from(Lifetime::Ping, store_name, None, &mut snapshotter);
137        storage.iter_store_from(Lifetime::Application, store_name, None, &mut snapshotter);
138        storage.iter_store_from(Lifetime::User, store_name, None, &mut snapshotter);
139
140        // Add send in all pings client.annotations
141        if store_name != "glean_client_info" {
142            storage.iter_store_from(Lifetime::Application, "all-pings", None, snapshotter);
143        }
144
145        if clear_store {
146            if let Err(e) = storage.clear_ping_lifetime_storage(store_name) {
147                log::warn!("Failed to clear lifetime storage: {:?}", e);
148            }
149        }
150
151        if snapshot.is_empty() {
152            None
153        } else {
154            Some(json!(snapshot))
155        }
156    }
157
158    /// Gets the current value of a single metric identified by name.
159    ///
160    /// # Arguments
161    ///
162    /// * `storage` - The database to get data from.
163    /// * `store_name` - The store name to look into.
164    /// * `metric_id` - The full metric identifier.
165    ///
166    /// # Returns
167    ///
168    /// The decoded metric or `None` if no data is found.
169    pub fn snapshot_metric(
170        &self,
171        storage: &Database,
172        store_name: &str,
173        metric_id: &str,
174        metric_lifetime: Lifetime,
175    ) -> Option<Metric> {
176        let mut snapshot: Option<Metric> = None;
177
178        let mut snapshotter = |id: &[u8], metric: &Metric| {
179            let id = String::from_utf8_lossy(id).into_owned();
180            if id == metric_id {
181                snapshot = Some(metric.clone())
182            }
183        };
184
185        storage.iter_store_from(metric_lifetime, store_name, None, &mut snapshotter);
186
187        snapshot
188    }
189
190    /// Gets the current value of a single metric identified by name.
191    ///
192    /// Use this API, rather than `snapshot_metric` within the testing API, so
193    /// that the usage will be reported in coverage, if enabled.
194    ///
195    /// # Arguments
196    ///
197    /// * `storage` - The database to get data from.
198    /// * `store_name` - The store name to look into.
199    /// * `metric_id` - The full metric identifier.
200    ///
201    /// # Returns
202    ///
203    /// The decoded metric or `None` if no data is found.
204    pub fn snapshot_metric_for_test(
205        &self,
206        storage: &Database,
207        store_name: &str,
208        metric_id: &str,
209        metric_lifetime: Lifetime,
210    ) -> Option<Metric> {
211        record_coverage(metric_id);
212        self.snapshot_metric(storage, store_name, metric_id, metric_lifetime)
213    }
214
215    ///  Snapshots the experiments.
216    ///
217    /// # Arguments
218    ///
219    /// * `storage` - The database to get data from.
220    /// * `store_name` - The store name to look into.
221    ///
222    /// # Returns
223    ///
224    /// A JSON representation of the experiment data, in the following format:
225    ///
226    /// ```json
227    /// {
228    ///  "experiment-id": {
229    ///    "branch": "branch-id",
230    ///    "extra": {
231    ///      "additional": "property",
232    ///      // ...
233    ///    }
234    ///  }
235    /// }
236    /// ```
237    ///
238    /// If no data for the store exists, `None` is returned.
239    pub fn snapshot_experiments_as_json(
240        &self,
241        storage: &Database,
242        store_name: &str,
243    ) -> Option<JsonValue> {
244        let mut snapshot: HashMap<String, JsonValue> = HashMap::new();
245
246        let mut snapshotter = |metric_id: &[u8], metric: &Metric| {
247            let metric_id = String::from_utf8_lossy(metric_id).into_owned();
248            if metric_id.ends_with("#experiment") {
249                let (name, _) = metric_id.split_once('#').unwrap(); // safe unwrap, we ensured there's a `#` in the string
250                snapshot.insert(name.to_string(), metric.as_json());
251            }
252        };
253
254        storage.iter_store_from(Lifetime::Application, store_name, None, &mut snapshotter);
255
256        if snapshot.is_empty() {
257            None
258        } else {
259            Some(json!(snapshot))
260        }
261    }
262}
263
264#[cfg(test)]
265mod test {
266    use super::*;
267    use crate::metrics::ExperimentMetric;
268    use crate::Glean;
269
270    // Experiment's API tests: the next test comes from glean-ac's
271    // ExperimentsStorageEngineTest.kt.
272    #[test]
273    fn test_experiments_json_serialization() {
274        let t = tempfile::tempdir().unwrap();
275        let name = t.path().display().to_string();
276        let glean = Glean::with_options(&name, "org.mozilla.glean", true, true);
277
278        let extra: HashMap<String, String> = [("test-key".into(), "test-value".into())]
279            .iter()
280            .cloned()
281            .collect();
282
283        let metric = ExperimentMetric::new(&glean, "some-experiment".to_string());
284
285        metric.set_active_sync(&glean, "test-branch".to_string(), extra);
286        let snapshot = StorageManager
287            .snapshot_experiments_as_json(glean.storage(), "glean_internal_info")
288            .unwrap();
289        assert_eq!(
290            json!({"some-experiment": {"branch": "test-branch", "extra": {"test-key": "test-value"}}}),
291            snapshot
292        );
293
294        metric.set_inactive_sync(&glean);
295
296        let empty_snapshot =
297            StorageManager.snapshot_experiments_as_json(glean.storage(), "glean_internal_info");
298        assert!(empty_snapshot.is_none());
299    }
300
301    #[test]
302    fn test_experiments_json_serialization_empty() {
303        let t = tempfile::tempdir().unwrap();
304        let name = t.path().display().to_string();
305        let glean = Glean::with_options(&name, "org.mozilla.glean", true, true);
306
307        let metric = ExperimentMetric::new(&glean, "some-experiment".to_string());
308
309        metric.set_active_sync(&glean, "test-branch".to_string(), HashMap::new());
310        let snapshot = StorageManager
311            .snapshot_experiments_as_json(glean.storage(), "glean_internal_info")
312            .unwrap();
313        assert_eq!(
314            json!({"some-experiment": {"branch": "test-branch"}}),
315            snapshot
316        );
317
318        metric.set_inactive_sync(&glean);
319
320        let empty_snapshot =
321            StorageManager.snapshot_experiments_as_json(glean.storage(), "glean_internal_info");
322        assert!(empty_snapshot.is_none());
323    }
324}
glean_core/storage/mod.rs

glean_core/storage/
mod.rs
