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

glean_core/storage/
mod.rs
