re_query/
storage_engine.rs

1use parking_lot::{ArcRwLockReadGuard, RwLockReadGuard};
2use re_chunk_store::{ChunkStore, ChunkStoreHandle};
3
4use crate::{QueryCache, QueryCacheHandle};
5
6// ---
7
8// TODO(cmc): This whole business should really be defined elsewhere, but for now this the best we
9// have, and it's really not worth adding yet another crate just for this.
10
11/// Anything that can expose references to a [`ChunkStore`] and its [`QueryCache`].
12///
13/// Used to abstract over [`StorageEngine`] and its different types of guards, such as [`StorageEngineArcReadGuard`].
14pub trait StorageEngineLike {
15    fn with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> R;
16
17    fn try_with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> Option<R> {
18        Some(self.with(f))
19    }
20}
21
22/// Keeps track of handles towards a [`ChunkStore`] and its [`QueryCache`].
23///
24/// A [`StorageEngine`] doesn't add any feature on top of what [`ChunkStoreHandle`] and
25/// [`QueryCacheHandle`] already offer: the job of the [`StorageEngine`] is to leverage the type
26/// system in order to protect against deadlocks and race conditions at compile time.
27///
28/// The handles stored within will never be publicly accessible past construction.
29///
30/// The underlying [`ChunkStore`] and [`QueryCache`] can be accessed through one of the
31/// following methods:
32/// * [`StorageEngine::read`]
33/// * [`StorageEngine::read_arc`]
34/// * [`StorageEngine::write`]
35/// * [`StorageEngine::write_arc`]
36#[derive(Clone)]
37pub struct StorageEngine {
38    store: ChunkStoreHandle,
39    cache: QueryCacheHandle,
40}
41
42impl StorageEngineLike for StorageEngine {
43    #[inline]
44    fn with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> R {
45        let this = self.read();
46        f(this.store(), this.cache())
47    }
48
49    #[inline]
50    fn try_with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> Option<R> {
51        let this = self.try_read()?;
52        Some(f(this.store(), this.cache()))
53    }
54}
55
56impl StorageEngine {
57    /// Creates a new [`StorageEngine`] with the specified [`ChunkStore`] and [`QueryCache`] handles.
58    ///
59    /// # Safety
60    ///
61    /// It is the responsibility of the caller to make sure that the given handles have not escaped
62    /// anywhere else before constructing this type, otherwise the [`StorageEngine`] cannot make
63    /// any safety guarantees.
64    #[inline]
65    #[expect(unsafe_code)]
66    pub unsafe fn new(store: ChunkStoreHandle, cache: QueryCacheHandle) -> Self {
67        Self { store, cache }
68    }
69}
70
71impl StorageEngine {
72    #[inline]
73    pub fn read(&self) -> StorageEngineReadGuard<'_> {
74        StorageEngineReadGuard {
75            cache: self.cache.read(),
76            store: self.store.read(),
77        }
78    }
79
80    #[inline]
81    pub fn try_read(&self) -> Option<StorageEngineReadGuard<'_>> {
82        let cache = self.cache.try_read()?;
83        let store = self.store.try_read()?;
84        Some(StorageEngineReadGuard { store, cache })
85    }
86
87    #[inline]
88    pub fn try_read_arc(&self) -> Option<StorageEngineArcReadGuard> {
89        let cache = self.cache.try_read_arc()?;
90        let store = self.store.try_read_arc()?;
91        Some(StorageEngineArcReadGuard { store, cache })
92    }
93
94    #[inline]
95    pub fn write(&self) -> StorageEngineWriteGuard<'_> {
96        StorageEngineWriteGuard {
97            cache: self.cache.write(),
98            store: self.store.write(),
99        }
100    }
101
102    #[inline]
103    pub fn try_write(&self) -> Option<StorageEngineWriteGuard<'_>> {
104        let cache = self.cache.try_write()?;
105        let store = self.store.try_write()?;
106        Some(StorageEngineWriteGuard { store, cache })
107    }
108
109    #[inline]
110    pub fn read_arc(&self) -> StorageEngineArcReadGuard {
111        StorageEngineArcReadGuard {
112            cache: self.cache.read_arc(),
113            store: self.store.read_arc(),
114        }
115    }
116
117    #[inline]
118    pub fn write_arc(&self) -> StorageEngineArcWriteGuard {
119        StorageEngineArcWriteGuard {
120            cache: self.cache.write_arc(),
121            store: self.store.write_arc(),
122        }
123    }
124
125    #[inline]
126    pub fn try_write_arc(&self) -> Option<StorageEngineArcWriteGuard> {
127        let cache = self.cache.try_write_arc()?;
128        let store = self.store.try_write_arc()?;
129        Some(StorageEngineArcWriteGuard { store, cache })
130    }
131}
132
133// --- Read Guards ---
134
135// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
136// as it is always possible to go back to an actual `RwLock` via `RwLockReadGuard::rwlock`.
137// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
138// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
139pub struct StorageEngineReadGuard<'a> {
140    store: parking_lot::RwLockReadGuard<'a, ChunkStore>,
141    cache: parking_lot::RwLockReadGuard<'a, QueryCache>,
142}
143
144impl Clone for StorageEngineReadGuard<'_> {
145    // Cloning the guard is safe, since the lock stays locked all along.
146    fn clone(&self) -> Self {
147        Self {
148            store: parking_lot::RwLock::read(RwLockReadGuard::rwlock(&self.store)),
149            cache: parking_lot::RwLock::read(RwLockReadGuard::rwlock(&self.cache)),
150        }
151    }
152}
153
154impl StorageEngineReadGuard<'_> {
155    #[inline]
156    pub fn store(&self) -> &ChunkStore {
157        &self.store
158    }
159
160    #[inline]
161    pub fn cache(&self) -> &QueryCache {
162        &self.cache
163    }
164}
165
166impl StorageEngineLike for StorageEngineReadGuard<'_> {
167    #[inline]
168    fn with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> R {
169        f(self.store(), self.cache())
170    }
171}
172
173// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
174// as it is always possible to go back to an actual `RwLock` via `ArcRwLockReadGuard::rwlock`.
175// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
176// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
177pub struct StorageEngineArcReadGuard {
178    store: parking_lot::ArcRwLockReadGuard<parking_lot::RawRwLock, ChunkStore>,
179    cache: parking_lot::ArcRwLockReadGuard<parking_lot::RawRwLock, QueryCache>,
180}
181
182impl StorageEngineArcReadGuard {
183    #[inline]
184    pub fn store(&self) -> &ChunkStore {
185        &self.store
186    }
187
188    #[inline]
189    pub fn cache(&self) -> &QueryCache {
190        &self.cache
191    }
192}
193
194impl StorageEngineLike for StorageEngineArcReadGuard {
195    #[inline]
196    fn with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> R {
197        f(self.store(), self.cache())
198    }
199}
200
201impl Clone for StorageEngineArcReadGuard {
202    // Cloning the guard is safe, since the lock stays locked all along.
203    fn clone(&self) -> Self {
204        Self {
205            store: parking_lot::RwLock::read_arc(ArcRwLockReadGuard::rwlock(&self.store)),
206            cache: parking_lot::RwLock::read_arc(ArcRwLockReadGuard::rwlock(&self.cache)),
207        }
208    }
209}
210
211// --- Write Guards ---
212
213// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
214// as it is always possible to go back to an actual `RwLock` via `RwLockWriteGuard::rwlock`.
215// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
216// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
217pub struct StorageEngineWriteGuard<'a> {
218    store: parking_lot::RwLockWriteGuard<'a, ChunkStore>,
219    cache: parking_lot::RwLockWriteGuard<'a, QueryCache>,
220}
221
222impl<'a> StorageEngineWriteGuard<'a> {
223    #[inline]
224    pub fn downgrade(self) -> StorageEngineReadGuard<'a> {
225        StorageEngineReadGuard {
226            store: parking_lot::RwLockWriteGuard::downgrade(self.store),
227            cache: parking_lot::RwLockWriteGuard::downgrade(self.cache),
228        }
229    }
230}
231
232impl StorageEngineWriteGuard<'_> {
233    #[inline]
234    pub fn store(&mut self) -> &mut ChunkStore {
235        &mut self.store
236    }
237
238    #[inline]
239    pub fn cache(&mut self) -> &mut QueryCache {
240        &mut self.cache
241    }
242}
243
244// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
245// as it is always possible to go back to an actual `RwLock` via `ArcRwLockWriteGuard::rwlock`.
246// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
247// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
248pub struct StorageEngineArcWriteGuard {
249    store: parking_lot::ArcRwLockWriteGuard<parking_lot::RawRwLock, ChunkStore>,
250    cache: parking_lot::ArcRwLockWriteGuard<parking_lot::RawRwLock, QueryCache>,
251}
252
253impl StorageEngineArcWriteGuard {
254    #[inline]
255    pub fn downgrade(self) -> StorageEngineArcReadGuard {
256        StorageEngineArcReadGuard {
257            store: parking_lot::ArcRwLockWriteGuard::downgrade(self.store),
258            cache: parking_lot::ArcRwLockWriteGuard::downgrade(self.cache),
259        }
260    }
261}
262
263impl StorageEngineArcWriteGuard {
264    #[inline]
265    pub fn store(&mut self) -> &mut ChunkStore {
266        &mut self.store
267    }
268
269    #[inline]
270    pub fn cache(&mut self) -> &mut QueryCache {
271        &mut self.cache
272    }
273}