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
31use objects::sync::LockExt;
32
33/// A process-lifetime cache that maps `K → V` and computes each entry on
34/// first access. See module docs for semantics.
35pub struct OnceMap<K, V> {
36    inner: OnceLock<Mutex<HashMap<K, V>>>,
37}
38
39impl<K, V> OnceMap<K, V> {
40    /// Empty cache, suitable for `static` initializers.
41    pub const fn new() -> Self {
42        Self {
43            inner: OnceLock::new(),
44        }
45    }
46
47    fn map(&self) -> &Mutex<HashMap<K, V>> {
48        self.inner.get_or_init(|| Mutex::new(HashMap::new()))
49    }
50}
51
52impl<K: Eq + Hash + Clone, V: Clone> OnceMap<K, V> {
53    /// Return the value for `key`, computing and caching it with `init`
54    /// on first access. The lock is held across `init`, so concurrent
55    /// callers for the same key serialize.
56    pub fn get_or_init_with<F>(&self, key: &K, init: F) -> V
57    where
58        F: FnOnce() -> V,
59    {
60        let mut guard = self.map().lock_or_poisoned();
61        if let Some(v) = guard.get(key) {
62            return v.clone();
63        }
64        let v = init();
65        guard.insert(key.clone(), v.clone());
66        v
67    }
68
69    /// Async variant of [`Self::get_or_init_with`]. The lock is released
70    /// across the await, so different keys don't serialize. Two callers
71    /// for the same key may both run `init`; the last write wins.
72    pub async fn get_or_init_async<F, Fut>(&self, key: &K, init: F) -> V
73    where
74        F: FnOnce() -> Fut,
75        Fut: std::future::Future<Output = V>,
76    {
77        if let Some(v) = self.get(key) {
78            return v;
79        }
80        let v = init().await;
81        self.map().lock_or_poisoned().insert(key.clone(), v.clone());
82        v
83    }
84
85    /// Read without computing. Returns `None` if the key was never inserted.
86    pub fn get(&self, key: &K) -> Option<V> {
87        self.map().lock_or_poisoned().get(key).cloned()
88    }
89
90    /// Direct insert. Returns the previous value if any.
91    pub fn insert(&self, key: K, value: V) -> Option<V> {
92        self.map().lock_or_poisoned().insert(key, value)
93    }
94
95    /// Remove and return the value for `key`, if any. Used by call
96    /// sites that need to tear down a cached resource (the mount
97    /// registry hands the handle back so the caller can unmount it).
98    pub fn remove(&self, key: &K) -> Option<V> {
99        self.map().lock_or_poisoned().remove(key)
100    }
101}
102
103impl<K, V> Default for OnceMap<K, V> {
104    fn default() -> Self {
105        Self::new()
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn computes_once_per_key() {
115        let map: OnceMap<String, i32> = OnceMap::new();
116        let mut calls = 0;
117        let mut init = |v: i32| {
118            calls += 1;
119            v
120        };
121        let a = map.get_or_init_with(&"a".to_string(), || init(1));
122        let a_again = map.get_or_init_with(&"a".to_string(), || init(2));
123        let b = map.get_or_init_with(&"b".to_string(), || init(3));
124        assert_eq!(a, 1);
125        assert_eq!(a_again, 1);
126        assert_eq!(b, 3);
127        assert_eq!(calls, 2);
128    }
129
130    #[test]
131    fn get_returns_none_when_missing() {
132        let map: OnceMap<String, i32> = OnceMap::new();
133        assert!(map.get(&"missing".to_string()).is_none());
134        map.insert("present".to_string(), 7);
135        assert_eq!(map.get(&"present".to_string()), Some(7));
136    }
137
138    #[tokio::test]
139    async fn async_init_caches_value() {
140        let map: OnceMap<String, i32> = OnceMap::new();
141        let v = map
142            .get_or_init_async(&"k".to_string(), || async { 42 })
143            .await;
144        assert_eq!(v, 42);
145        assert_eq!(map.get(&"k".to_string()), Some(42));
146    }
147}
cli/util/once_map.rs

cli/util/
once_map.rs
