vortex_session/
arc_swap_map.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4//! A concurrent, copy-on-write map backed by an [`ArcSwap`].
5
6use std::borrow::Borrow;
7use std::fmt::Debug;
8use std::fmt::Formatter;
9use std::hash::BuildHasher;
10use std::hash::Hash;
11use std::sync::Arc;
12
13use arc_swap::ArcSwap;
14use arc_swap::Guard;
15use vortex_utils::aliases::hash_map::DefaultHashBuilder;
16use vortex_utils::aliases::hash_map::HashMap;
17
18/// A concurrent [`HashMap`] backed by an [`ArcSwap`], offering lock-free reads
19/// and copy-on-write writes.
20///
21/// Reads load the current snapshot without blocking writers. Writes clone the
22/// whole map, apply their change, and atomically publish the new version, so a
23/// reader always observes a consistent snapshot and writers never block readers.
24///
25/// This is the shared building block behind the session-scoped registries (the
26/// plugin registries as well as the optimizer-kernel and aggregate-function
27/// registries) and the [`VortexSession`](crate::VortexSession) type-map itself.
28/// Because every write clones the entire map, it is intended for maps that are
29/// written rarely (typically only while a session is being configured) and read
30/// often.
31///
32/// The map is held behind an [`Arc`] so that [`Clone`] shares the same
33/// underlying cell: a registry mutated through one clone is observed by all
34/// others. Session variables rely on this so that encodings registered after a
35/// session is built remain visible to clones of that session.
36pub struct ArcSwapMap<K, V, S = DefaultHashBuilder> {
37    inner: Arc<ArcSwap<HashMap<K, V, S>>>,
38}
39
40impl<K, V, S: Default> Default for ArcSwapMap<K, V, S> {
41    fn default() -> Self {
42        Self {
43            inner: Arc::new(ArcSwap::from_pointee(HashMap::default())),
44        }
45    }
46}
47
48impl<K, V, S> Clone for ArcSwapMap<K, V, S> {
49    fn clone(&self) -> Self {
50        Self {
51            inner: Arc::clone(&self.inner),
52        }
53    }
54}
55
56impl<K: Debug, V: Debug, S> Debug for ArcSwapMap<K, V, S> {
57    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
58        self.read(|map| f.debug_tuple("ArcSwapMap").field(map).finish())
59    }
60}
61
62impl<K, V, S> ArcSwapMap<K, V, S> {
63    /// Return the currently published map snapshot.
64    pub fn snapshot(&self) -> Arc<HashMap<K, V, S>> {
65        self.inner.load_full()
66    }
67
68    /// Read the current snapshot, passing it to `f`.
69    ///
70    /// Every lookup inside `f` observes the same snapshot, which matters when a
71    /// single logical read consults more than one key.
72    pub fn read<R>(&self, f: impl FnOnce(&HashMap<K, V, S>) -> R) -> R {
73        f(&self.inner.load())
74    }
75
76    /// Return a lock-free guard to the current snapshot without cloning the
77    /// [`Arc`].
78    ///
79    /// Cheaper than [`snapshot`](Self::snapshot) for short-lived reads, but the
80    /// guard pins an internal arc-swap slot, so it should stay on the stack
81    /// rather than be stored in a long-lived data structure.
82    pub(crate) fn load(&self) -> Guard<Arc<HashMap<K, V, S>>> {
83        self.inner.load()
84    }
85
86    /// Replace the map with the result of applying `f` to a private copy.
87    ///
88    /// Writes are copy-on-write via [`ArcSwap::rcu`], so `f` may run more than
89    /// once under contention and must not move out of its captures.
90    fn modify(&self, f: impl Fn(&mut HashMap<K, V, S>))
91    where
92        K: Clone,
93        V: Clone,
94        S: Clone,
95    {
96        self.inner.rcu(|existing| {
97            let mut map = existing.as_ref().clone();
98            f(&mut map);
99            map
100        });
101    }
102}
103
104impl<K: Eq + Hash, V, S: BuildHasher> ArcSwapMap<K, V, S> {
105    /// Return a clone of the value stored under `key`, if present.
106    pub fn get<Q>(&self, key: &Q) -> Option<V>
107    where
108        K: Borrow<Q>,
109        Q: Eq + Hash + ?Sized,
110        V: Clone,
111    {
112        self.inner.load().get(key).cloned()
113    }
114
115    /// Returns whether a value is stored under `key`.
116    pub fn contains_key<Q>(&self, key: &Q) -> bool
117    where
118        K: Borrow<Q>,
119        Q: Eq + Hash + ?Sized,
120    {
121        self.inner.load().contains_key(key)
122    }
123
124    pub fn insert(&self, key: K, value: V)
125    where
126        K: Clone,
127        V: Clone,
128        S: Clone,
129    {
130        self.modify(|map| {
131            map.insert(key.clone(), value.clone());
132        });
133    }
134
135    /// Insert `value` under `key` only if no value is stored there yet.
136    ///
137    /// If a concurrent writer publishes a value under `key` first, that value
138    /// is kept and `value` is dropped. `value` is constructed by the caller
139    /// before this call, so no user code runs while the map is being updated.
140    pub fn insert_if_absent(&self, key: K, value: V)
141    where
142        K: Clone,
143        V: Clone,
144        S: Clone,
145    {
146        self.modify(|map| {
147            map.entry(key.clone()).or_insert_with(|| value.clone());
148        });
149    }
150}
151
152impl<K: Eq + Hash + Clone, T: Clone, S: BuildHasher + Clone> ArcSwapMap<K, Arc<[T]>, S> {
153    /// Append `values` to the list stored under `key`, creating it if absent.
154    ///
155    /// Each key maps to an immutable `Arc<[T]>`; appending rebuilds that slice
156    /// copy-on-write so existing readers keep their previous snapshot.
157    pub fn extend(&self, key: K, values: &[T]) {
158        self.modify(|map| {
159            let merged: Arc<[T]> = match map.get(&key) {
160                Some(existing) => existing.iter().chain(values).cloned().collect(),
161                None => values.into(),
162            };
163            map.insert(key.clone(), merged);
164        });
165    }
166
167    /// Append a single `value` to the list stored under `key`, creating it if
168    /// absent.
169    pub fn push(&self, key: K, value: T) {
170        self.extend(key, &[value]);
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn get_and_insert() {
180        let map = ArcSwapMap::<u32, i32>::default();
181        assert_eq!(map.get(&1), None);
182        assert!(!map.contains_key(&1));
183        map.insert(1, 10);
184        map.insert(1, 20);
185        assert_eq!(map.get(&1), Some(20));
186        assert!(map.contains_key(&1));
187    }
188
189    #[test]
190    fn insert_if_absent_keeps_first_value() {
191        let map = ArcSwapMap::<u32, i32>::default();
192        map.insert_if_absent(1, 10);
193        map.insert_if_absent(1, 20);
194        assert_eq!(map.get(&1), Some(10));
195    }
196
197    #[test]
198    fn extend_appends_per_key() {
199        let map = ArcSwapMap::<u32, Arc<[i32]>>::default();
200        map.extend(1, &[1, 2]);
201        map.extend(1, &[3]);
202        map.extend(2, &[4]);
203        assert_eq!(map.get(&1).as_deref(), Some([1, 2, 3].as_slice()));
204        assert_eq!(map.get(&2).as_deref(), Some([4].as_slice()));
205    }
206
207    #[test]
208    fn push_appends_single_values() {
209        let map = ArcSwapMap::<u32, Arc<[i32]>>::default();
210        map.push(1, 1);
211        map.push(1, 2);
212        assert_eq!(map.get(&1).as_deref(), Some([1, 2].as_slice()));
213    }
214
215    #[test]
216    fn read_observes_a_single_snapshot() {
217        let map = ArcSwapMap::<u32, i32>::default();
218        map.insert(1, 1);
219        map.insert(2, 2);
220        assert_eq!(map.read(|m| m.values().sum::<i32>()), 3);
221    }
222
223    #[test]
224    fn snapshot_keeps_published_view() {
225        let map = ArcSwapMap::<u32, i32>::default();
226        map.insert(1, 10);
227
228        let snapshot = map.snapshot();
229        map.insert(1, 20);
230        map.insert(2, 30);
231
232        assert_eq!(snapshot.get(&1), Some(&10));
233        assert_eq!(snapshot.get(&2), None);
234        assert_eq!(map.get(&1), Some(20));
235        assert_eq!(map.get(&2), Some(30));
236    }
237
238    #[test]
239    fn clone_shares_the_same_cell() {
240        let map = ArcSwapMap::<u32, i32>::default();
241        let clone = map.clone();
242        // A write through one handle is observed through the other.
243        map.insert(1, 10);
244        assert_eq!(clone.get(&1), Some(10));
245        clone.insert(2, 20);
246        assert_eq!(map.get(&2), Some(20));
247    }
248}
vortex_session/arc_swap_map.rs

vortex_session/
arc_swap_map.rs
