nym_cache/
lib.rs

1// Copyright 2025 - Nym Technologies SA <contact@nymtech.net>
2// SPDX-License-Identifier: Apache-2.0
3
4use std::collections::HashMap;
5use std::future::Future;
6use std::hash::Hash;
7use std::ops::Deref;
8use tokio::sync::{RwLock, RwLockReadGuard};
9
10/// A map of items that never change for given key
11pub struct CachedImmutableItems<K, V> {
12    // I wonder if there's a more efficient structure with OnceLock or OnceCell or something
13    inner: RwLock<HashMap<K, V>>,
14}
15
16impl<K, V> Default for CachedImmutableItems<K, V> {
17    fn default() -> Self {
18        CachedImmutableItems {
19            inner: RwLock::new(HashMap::new()),
20        }
21    }
22}
23
24impl<K, V> Deref for CachedImmutableItems<K, V> {
25    type Target = RwLock<HashMap<K, V>>;
26
27    fn deref(&self) -> &Self::Target {
28        &self.inner
29    }
30}
31
32impl<K, V> CachedImmutableItems<K, V>
33where
34    K: Eq + Hash,
35{
36    pub async fn get_or_init<F, U, E>(
37        &self,
38        key: K,
39        initialiser: F,
40    ) -> Result<RwLockReadGuard<'_, V>, E>
41    where
42        F: FnOnce() -> U,
43        U: Future<Output = Result<V, E>>,
44        K: Clone,
45    {
46        // 1. see if we already have the item cached
47        let guard = self.inner.read().await;
48        if let Ok(item) = RwLockReadGuard::try_map(guard, |map| map.get(&key)) {
49            return Ok(item);
50        }
51
52        // 2. attempt to retrieve (and cache) it
53        let mut write_guard = self.inner.write().await;
54
55        // see if another task has already set the item whilst we were waiting for the lock
56        if write_guard.get(&key).is_some() {
57            let read_guard = write_guard.downgrade();
58
59            // SAFETY: we just checked the entry exists and we never dropped the guard
60            #[allow(clippy::unwrap_used)]
61            return Ok(RwLockReadGuard::map(read_guard, |map| {
62                map.get(&key).unwrap()
63            }));
64        }
65
66        let init = initialiser().await?;
67        write_guard.insert(key.clone(), init);
68
69        let guard = write_guard.downgrade();
70
71        // SAFETY:
72        // we just inserted the entry into the map while NEVER dropping the lock (only downgraded it)
73        // so it MUST exist and thus the unwrap is fine
74        #[allow(clippy::unwrap_used)]
75        Ok(RwLockReadGuard::map(guard, |map| map.get(&key).unwrap()))
76    }
77}