vortex_session/
lib.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4pub mod registry;
5
6use std::any::Any;
7use std::any::TypeId;
8use std::any::type_name;
9use std::fmt::Debug;
10use std::hash::BuildHasherDefault;
11use std::hash::Hasher;
12use std::ops::Deref;
13use std::ops::DerefMut;
14use std::sync::Arc;
15
16use dashmap::DashMap;
17use dashmap::Entry;
18use vortex_error::VortexExpect;
19use vortex_error::vortex_panic;
20
21/// A Vortex session encapsulates the set of extensible arrays, layouts, compute functions, dtypes,
22/// etc. that are available for use in a given context.
23///
24/// It is also the entry-point passed to dynamic libraries to initialize Vortex plugins.
25#[derive(Clone, Debug)]
26pub struct VortexSession(Arc<SessionVars>);
27
28impl VortexSession {
29    /// Create a new [`VortexSession`] with no session state.
30    ///
31    /// It is recommended to use the `default()` method instead provided by the main `vortex` crate.
32    pub fn empty() -> Self {
33        Self(Default::default())
34    }
35
36    /// Inserts a new session variable of type `V` with its default value.
37    ///
38    /// # Panics
39    ///
40    /// If a variable of that type already exists.
41    pub fn with<V: SessionVar + Default>(self) -> Self {
42        match self.0.entry(TypeId::of::<V>()) {
43            Entry::Occupied(_) => {
44                vortex_panic!(
45                    "Session variable of type {} already exists",
46                    type_name::<V>()
47                );
48            }
49            Entry::Vacant(e) => {
50                e.insert(Box::new(V::default()));
51            }
52        }
53        self
54    }
55}
56
57/// Trait for accessing and modifying the state of a Vortex session.
58pub trait SessionExt: Sized + private::Sealed {
59    /// Returns the [`VortexSession`].
60    fn session(&self) -> VortexSession;
61
62    /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
63    fn get<V: SessionVar + Default>(&self) -> Ref<'_, V>;
64
65    /// Returns the scope variable of type `V` if it exists.
66    fn get_opt<V: SessionVar>(&self) -> Option<Ref<'_, V>>;
67
68    /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
69    ///
70    /// Note that the returned value internally holds a lock on the variable.
71    fn get_mut<V: SessionVar + Default>(&self) -> RefMut<'_, V>;
72
73    /// Returns the scope variable of type `V`, if it exists.
74    ///
75    /// Note that the returned value internally holds a lock on the variable.
76    fn get_mut_opt<V: SessionVar>(&self) -> Option<RefMut<'_, V>>;
77}
78
79mod private {
80    pub trait Sealed {}
81    impl Sealed for super::VortexSession {}
82}
83
84impl SessionExt for VortexSession {
85    fn session(&self) -> VortexSession {
86        self.clone()
87    }
88
89    /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
90    fn get<V: SessionVar + Default>(&self) -> Ref<'_, V> {
91        // NOTE(ngates): we don't use `entry().or_insert_with_key()` here because the DashMap
92        //  would immediately acquire an exclusive write lock.
93        if let Some(v) = self.0.get(&TypeId::of::<V>()) {
94            return Ref(v.map(|v| {
95                (**v)
96                    .as_any()
97                    .downcast_ref::<V>()
98                    .vortex_expect("Type mismatch - this is a bug")
99            }));
100        }
101
102        // If we get here, the value was not present, so we insert the default with a write lock.
103        Ref(self
104            .0
105            .entry(TypeId::of::<V>())
106            .or_insert_with(|| Box::new(V::default()))
107            .downgrade()
108            .map(|v| {
109                (**v)
110                    .as_any()
111                    .downcast_ref::<V>()
112                    .vortex_expect("Type mismatch - this is a bug")
113            }))
114    }
115
116    fn get_opt<V: SessionVar>(&self) -> Option<Ref<'_, V>> {
117        self.0.get(&TypeId::of::<V>()).map(|v| {
118            Ref(v.map(|v| {
119                (**v)
120                    .as_any()
121                    .downcast_ref::<V>()
122                    .vortex_expect("Type mismatch - this is a bug")
123            }))
124        })
125    }
126
127    /// Returns the scope variable of type `V`, or inserts a default one if it does not exist.
128    ///
129    /// Note that the returned value internally holds a lock on the variable.
130    fn get_mut<V: SessionVar + Default>(&self) -> RefMut<'_, V> {
131        RefMut(
132            self.0
133                .entry(TypeId::of::<V>())
134                .or_insert_with(|| Box::new(V::default()))
135                .map(|v| {
136                    (**v)
137                        .as_any_mut()
138                        .downcast_mut::<V>()
139                        .vortex_expect("Type mismatch - this is a bug")
140                }),
141        )
142    }
143
144    fn get_mut_opt<V: SessionVar>(&self) -> Option<RefMut<'_, V>> {
145        self.0.get_mut(&TypeId::of::<V>()).map(|v| {
146            RefMut(v.map(|v| {
147                (**v)
148                    .as_any_mut()
149                    .downcast_mut::<V>()
150                    .vortex_expect("Type mismatch - this is a bug")
151            }))
152        })
153    }
154}
155
156/// A TypeMap based on `https://docs.rs/http/1.2.0/src/http/extensions.rs.html#41-266`.
157type SessionVars = DashMap<TypeId, Box<dyn SessionVar>, BuildHasherDefault<IdHasher>>;
158
159/// With TypeIds as keys, there's no need to hash them. They are already hashes
160/// themselves, coming from the compiler. The IdHasher just holds the u64 of
161/// the TypeId, and then returns it, instead of doing any bit fiddling.
162#[derive(Default)]
163struct IdHasher(u64);
164
165impl Hasher for IdHasher {
166    #[inline]
167    fn finish(&self) -> u64 {
168        self.0
169    }
170
171    fn write(&mut self, _: &[u8]) {
172        unreachable!("TypeId calls write_u64");
173    }
174
175    #[inline]
176    fn write_u64(&mut self, id: u64) {
177        self.0 = id;
178    }
179}
180
181/// This trait defines variables that can be stored against a Vortex session.
182pub trait SessionVar: Any + Send + Sync + Debug + 'static {
183    fn as_any(&self) -> &dyn Any;
184    fn as_any_mut(&mut self) -> &mut dyn Any;
185}
186
187impl<T: Send + Sync + Debug + 'static> SessionVar for T {
188    fn as_any(&self) -> &dyn Any {
189        self
190    }
191
192    fn as_any_mut(&mut self) -> &mut dyn Any {
193        self
194    }
195}
196
197// NOTE(ngates): we don't want to expose that the internals of a session is a DashMap, so we have
198// our own wrapped Ref type.
199pub struct Ref<'a, T>(dashmap::mapref::one::MappedRef<'a, TypeId, Box<dyn SessionVar>, T>);
200impl<'a, T> Deref for Ref<'a, T> {
201    type Target = T;
202
203    fn deref(&self) -> &Self::Target {
204        &self.0
205    }
206}
207impl<'a, T> Ref<'a, T> {
208    /// Map this reference to a different target.
209    pub fn map<F, U>(self, f: F) -> Ref<'a, U>
210    where
211        F: FnOnce(&T) -> &U,
212    {
213        Ref(self.0.map(f))
214    }
215}
216
217pub struct RefMut<'a, T>(dashmap::mapref::one::MappedRefMut<'a, TypeId, Box<dyn SessionVar>, T>);
218impl<'a, T> Deref for RefMut<'a, T> {
219    type Target = T;
220
221    fn deref(&self) -> &Self::Target {
222        &self.0
223    }
224}
225impl<'a, T> DerefMut for RefMut<'a, T> {
226    fn deref_mut(&mut self) -> &mut Self::Target {
227        self.0.deref_mut()
228    }
229}
230impl<'a, T> RefMut<'a, T> {
231    /// Map this mutable reference to a different target.
232    pub fn map<F, U>(self, f: F) -> RefMut<'a, U>
233    where
234        F: FnOnce(&mut T) -> &mut U,
235    {
236        RefMut(self.0.map(f))
237    }
238}
vortex_session/lib.rs

vortex_session/
lib.rs
