Skip to main content

metrics_runtime/registry/
scope.rs

1use crate::common::{Scope, ScopeHandle};
2use parking_lot::RwLock;
3use std::collections::HashMap;
4
5#[derive(Debug)]
6struct Inner {
7    id: u64,
8    forward: HashMap<Scope, ScopeHandle>,
9    backward: HashMap<ScopeHandle, Scope>,
10}
11
12impl Inner {
13    pub fn new() -> Self {
14        Inner {
15            id: 1,
16            forward: HashMap::new(),
17            backward: HashMap::new(),
18        }
19    }
20}
21
22#[derive(Debug)]
23pub(crate) struct ScopeRegistry {
24    inner: RwLock<Inner>,
25}
26
27impl ScopeRegistry {
28    pub fn new() -> Self {
29        Self {
30            inner: RwLock::new(Inner::new()),
31        }
32    }
33
34    pub fn register(&self, scope: Scope) -> u64 {
35        let mut wg = self.inner.write();
36
37        // If the key is already registered, send back the existing scope ID.
38        if wg.forward.contains_key(&scope) {
39            return wg.forward.get(&scope).cloned().unwrap();
40        }
41
42        // Otherwise, take the current scope ID for this registration, store it, and increment
43        // the scope ID counter for the next registration.
44        let scope_id = wg.id;
45        let _ = wg.forward.insert(scope.clone(), scope_id);
46        let _ = wg.backward.insert(scope_id, scope);
47        wg.id += 1;
48        scope_id
49    }
50
51    pub fn get(&self, scope_id: ScopeHandle) -> Scope {
52        // See if we have an entry for the scope ID, and clone the scope if so.
53        let rg = self.inner.read();
54        rg.backward.get(&scope_id).cloned().unwrap_or(Scope::Root)
55    }
56}
57
58#[cfg(test)]
59mod tests {
60    use super::{Scope, ScopeRegistry};
61
62    #[test]
63    fn test_simple_write_then_read() {
64        let nested1 = Scope::Root.add_part("nested1");
65        let nested2 = nested1.clone().add_part("nested2");
66
67        let sr = ScopeRegistry::new();
68
69        let doesnt_exist0 = sr.get(0);
70        let doesnt_exist1 = sr.get(1);
71        let doesnt_exist2 = sr.get(2);
72
73        assert_eq!(doesnt_exist0, Scope::Root);
74        assert_eq!(doesnt_exist1, Scope::Root);
75        assert_eq!(doesnt_exist2, Scope::Root);
76
77        let nested1_original = nested1.clone();
78        let nested1_id = sr.register(nested1);
79
80        let nested2_original = nested2.clone();
81        let nested2_id = sr.register(nested2);
82
83        let exists1 = sr.get(nested1_id);
84        let exists2 = sr.get(nested2_id);
85
86        assert_eq!(exists1, nested1_original);
87        assert_eq!(exists2, nested2_original);
88    }
89}