uv_once_map/
lib.rs

1use std::borrow::Borrow;
2use std::fmt::{Debug, Formatter};
3use std::hash::{BuildHasher, Hash, RandomState};
4use std::sync::Arc;
5
6use dashmap::DashMap;
7use tokio::sync::Notify;
8
9/// Run tasks only once and store the results in a parallel hash map.
10///
11/// We often have jobs `Fn(K) -> V` that we only want to run once and memoize, e.g. network
12/// requests for metadata. When multiple tasks start the same query in parallel, e.g. through source
13/// dist builds, we want to wait until the other task is done and get a reference to the same
14/// result.
15///
16/// Note that this always clones the value out of the underlying map. Because
17/// of this, it's common to wrap the `V` in an `Arc<V>` to make cloning cheap.
18pub struct OnceMap<K, V, S = RandomState> {
19    items: DashMap<K, Value<V>, S>,
20}
21
22impl<K: Eq + Hash + Debug, V: Debug, S: BuildHasher + Clone> Debug for OnceMap<K, V, S> {
23    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
24        Debug::fmt(&self.items, f)
25    }
26}
27
28impl<K: Eq + Hash, V: Clone, H: BuildHasher + Clone> OnceMap<K, V, H> {
29    /// Create a [`OnceMap`] with the specified hasher.
30    pub fn with_hasher(hasher: H) -> Self {
31        Self {
32            items: DashMap::with_hasher(hasher),
33        }
34    }
35
36    /// Create a [`OnceMap`] with the specified capacity and hasher.
37    pub fn with_capacity_and_hasher(capacity: usize, hasher: H) -> Self {
38        Self {
39            items: DashMap::with_capacity_and_hasher(capacity, hasher),
40        }
41    }
42
43    /// Register that you want to start a job.
44    ///
45    /// If this method returns `true`, you need to start a job and call [`OnceMap::done`] eventually
46    /// or other tasks will hang. If it returns `false`, this job is already in progress and you
47    /// can [`OnceMap::wait`] for the result.
48    pub fn register(&self, key: K) -> bool {
49        let entry = self.items.entry(key);
50        match entry {
51            dashmap::mapref::entry::Entry::Occupied(_) => false,
52            dashmap::mapref::entry::Entry::Vacant(entry) => {
53                entry.insert(Value::Waiting(Arc::new(Notify::new())));
54                true
55            }
56        }
57    }
58
59    /// Submit the result of a job you registered.
60    pub fn done(&self, key: K, value: V) {
61        if let Some(Value::Waiting(notify)) = self.items.insert(key, Value::Filled(value)) {
62            notify.notify_waiters();
63        }
64    }
65
66    /// Wait for the result of a job that is running.
67    ///
68    /// Will hang if [`OnceMap::done`] isn't called for this key.
69    pub async fn wait(&self, key: &K) -> Option<V> {
70        let notify = {
71            let entry = self.items.get(key)?;
72            match entry.value() {
73                Value::Filled(value) => return Some(value.clone()),
74                Value::Waiting(notify) => notify.clone(),
75            }
76        };
77
78        // Register the waiter for calls to `notify_waiters`.
79        let notification = notify.notified();
80
81        // Make sure the value wasn't inserted in-between us checking the map and registering the waiter.
82        if let Value::Filled(value) = self.items.get(key).expect("map is append-only").value() {
83            return Some(value.clone());
84        }
85
86        // Wait until the value is inserted.
87        notification.await;
88
89        let entry = self.items.get(key).expect("map is append-only");
90        match entry.value() {
91            Value::Filled(value) => Some(value.clone()),
92            Value::Waiting(_) => unreachable!("notify was called"),
93        }
94    }
95
96    /// Wait for the result of a job that is running, in a blocking context.
97    ///
98    /// Will hang if [`OnceMap::done`] isn't called for this key.
99    pub fn wait_blocking(&self, key: &K) -> Option<V> {
100        futures::executor::block_on(self.wait(key))
101    }
102
103    /// Return the result of a previous job, if any.
104    pub fn get<Q: ?Sized + Hash + Eq>(&self, key: &Q) -> Option<V>
105    where
106        K: Borrow<Q>,
107    {
108        let entry = self.items.get(key)?;
109        match entry.value() {
110            Value::Filled(value) => Some(value.clone()),
111            Value::Waiting(_) => None,
112        }
113    }
114
115    /// Remove the result of a previous job, if any.
116    pub fn remove<Q: ?Sized + Hash + Eq>(&self, key: &Q) -> Option<V>
117    where
118        K: Borrow<Q>,
119    {
120        let entry = self.items.remove(key)?;
121        match entry {
122            (_, Value::Filled(value)) => Some(value),
123            (_, Value::Waiting(_)) => None,
124        }
125    }
126}
127
128impl<K: Eq + Hash + Clone, V, H: Default + BuildHasher + Clone> Default for OnceMap<K, V, H> {
129    fn default() -> Self {
130        Self {
131            items: DashMap::with_hasher(H::default()),
132        }
133    }
134}
135
136impl<K, V, H> FromIterator<(K, V)> for OnceMap<K, V, H>
137where
138    K: Eq + Hash,
139    H: Default + Clone + BuildHasher,
140{
141    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
142        Self {
143            items: iter
144                .into_iter()
145                .map(|(k, v)| (k, Value::Filled(v)))
146                .collect(),
147        }
148    }
149}
150
151#[derive(Debug)]
152enum Value<V> {
153    Waiting(Arc<Notify>),
154    Filled(V),
155}
uv_once_map/lib.rs

uv_once_map/
lib.rs
