cli/util/
once_map.rs

1// SPDX-License-Identifier: Apache-2.0
2//! Process-lifetime per-key cache.
3//!
4//! `OnceMap<K, V>` is the shape that the CLI's daemon detection and mount
5//! lifecycle code all reach for independently: a process-wide `HashMap`
6//! whose entries are computed on first access and held until the process
7//! exits. The repeated `OnceLock<Mutex<HashMap<K, V>>>` ceremony at each
8//! call site disappears behind a single type.
9//!
10//! Two construction modes:
11//!
12//! - [`OnceMap::get_or_init_with`] for synchronous initializers (file-stat
13//!   probes, key derivation).
14//! - [`OnceMap::get_or_init_async`] for async initializers (gRPC channel
15//!   construction, network handshakes). Concurrent inserts for *different*
16//!   keys don't serialize because the lock is released across the `await`;
17//!   concurrent inserts for the *same* key may both run the init future
18//!   (last writer wins) — this matches the behavior of every call site
19//!   that previously used this pattern.
20//!
21//! Values are cloned on read. Use a cheaply-cloneable handle (`Arc<…>`,
22//! `tonic::transport::Channel`) when the underlying object is expensive
23//! to clone.
24
25use std::{
26    collections::HashMap,
27    hash::Hash,
28    sync::{Mutex, OnceLock},
29};
30
31/// A process-lifetime cache that maps `K → V` and computes each entry on
32/// first access. See module docs for semantics.
33pub struct OnceMap<K, V> {
34    inner: OnceLock<Mutex<HashMap<K, V>>>,
35}
36
37impl<K, V> OnceMap<K, V> {
38    /// Empty cache, suitable for `static` initializers.
39    pub const fn new() -> Self {
40        Self {
41            inner: OnceLock::new(),
42        }
43    }
44
45    fn map(&self) -> &Mutex<HashMap<K, V>> {
46        self.inner.get_or_init(|| Mutex::new(HashMap::new()))
47    }
48}
49
50impl<K: Eq + Hash + Clone, V: Clone> OnceMap<K, V> {
51    /// Return the value for `key`, computing and caching it with `init`
52    /// on first access. The lock is held across `init`, so concurrent
53    /// callers for the same key serialize.
54    pub fn get_or_init_with<F>(&self, key: &K, init: F) -> V
55    where
56        F: FnOnce() -> V,
57    {
58        let mut guard = self.map().lock().expect("OnceMap mutex poisoned");
59        if let Some(v) = guard.get(key) {
60            return v.clone();
61        }
62        let v = init();
63        guard.insert(key.clone(), v.clone());
64        v
65    }
66
67    /// Async variant of [`Self::get_or_init_with`]. The lock is released
68    /// across the await, so different keys don't serialize. Two callers
69    /// for the same key may both run `init`; the last write wins.
70    pub async fn get_or_init_async<F, Fut>(&self, key: &K, init: F) -> V
71    where
72        F: FnOnce() -> Fut,
73        Fut: std::future::Future<Output = V>,
74    {
75        if let Some(v) = self.get(key) {
76            return v;
77        }
78        let v = init().await;
79        self.map()
80            .lock()
81            .expect("OnceMap mutex poisoned")
82            .insert(key.clone(), v.clone());
83        v
84    }
85
86    /// Read without computing. Returns `None` if the key was never inserted.
87    pub fn get(&self, key: &K) -> Option<V> {
88        self.map()
89            .lock()
90            .expect("OnceMap mutex poisoned")
91            .get(key)
92            .cloned()
93    }
94
95    /// Direct insert. Returns the previous value if any.
96    pub fn insert(&self, key: K, value: V) -> Option<V> {
97        self.map()
98            .lock()
99            .expect("OnceMap mutex poisoned")
100            .insert(key, value)
101    }
102
103    /// Remove and return the value for `key`, if any. Used by call
104    /// sites that need to tear down a cached resource (the mount
105    /// registry hands the handle back so the caller can unmount it).
106    pub fn remove(&self, key: &K) -> Option<V> {
107        self.map()
108            .lock()
109            .expect("OnceMap mutex poisoned")
110            .remove(key)
111    }
112}
113
114impl<K, V> Default for OnceMap<K, V> {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120#[cfg(test)]
121mod tests {
122    use super::*;
123
124    #[test]
125    fn computes_once_per_key() {
126        let map: OnceMap<String, i32> = OnceMap::new();
127        let mut calls = 0;
128        let mut init = |v: i32| {
129            calls += 1;
130            v
131        };
132        let a = map.get_or_init_with(&"a".to_string(), || init(1));
133        let a_again = map.get_or_init_with(&"a".to_string(), || init(2));
134        let b = map.get_or_init_with(&"b".to_string(), || init(3));
135        assert_eq!(a, 1);
136        assert_eq!(a_again, 1);
137        assert_eq!(b, 3);
138        assert_eq!(calls, 2);
139    }
140
141    #[test]
142    fn get_returns_none_when_missing() {
143        let map: OnceMap<String, i32> = OnceMap::new();
144        assert!(map.get(&"missing".to_string()).is_none());
145        map.insert("present".to_string(), 7);
146        assert_eq!(map.get(&"present".to_string()), Some(7));
147    }
148
149    #[tokio::test]
150    async fn async_init_caches_value() {
151        let map: OnceMap<String, i32> = OnceMap::new();
152        let v = map
153            .get_or_init_async(&"k".to_string(), || async { 42 })
154            .await;
155        assert_eq!(v, 42);
156        assert_eq!(map.get(&"k".to_string()), Some(42));
157    }
158}
cli/util/once_map.rs

cli/util/
once_map.rs
