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
173impl re_byte_size::MemUsageTreeCapture for StorageEngineReadGuard<'_> {
174    fn capture_mem_usage_tree(&self) -> re_byte_size::MemUsageTree {
175        let Self { store, cache } = self;
176        re_byte_size::MemUsageNode::new()
177            .with_child("ChunkStore", store.capture_mem_usage_tree())
178            .with_child("QueryCache", cache.capture_mem_usage_tree())
179            .into_tree()
180    }
181}
182
183// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
184// as it is always possible to go back to an actual `RwLock` via `ArcRwLockReadGuard::rwlock`.
185// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
186// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
187pub struct StorageEngineArcReadGuard {
188    store: parking_lot::ArcRwLockReadGuard<parking_lot::RawRwLock, ChunkStore>,
189    cache: parking_lot::ArcRwLockReadGuard<parking_lot::RawRwLock, QueryCache>,
190}
191
192impl StorageEngineArcReadGuard {
193    #[inline]
194    pub fn store(&self) -> &ChunkStore {
195        &self.store
196    }
197
198    #[inline]
199    pub fn cache(&self) -> &QueryCache {
200        &self.cache
201    }
202}
203
204impl StorageEngineLike for StorageEngineArcReadGuard {
205    #[inline]
206    fn with<F: FnOnce(&ChunkStore, &QueryCache) -> R, R>(&self, f: F) -> R {
207        f(self.store(), self.cache())
208    }
209}
210
211impl Clone for StorageEngineArcReadGuard {
212    // Cloning the guard is safe, since the lock stays locked all along.
213    fn clone(&self) -> Self {
214        Self {
215            store: parking_lot::RwLock::read_arc(ArcRwLockReadGuard::rwlock(&self.store)),
216            cache: parking_lot::RwLock::read_arc(ArcRwLockReadGuard::rwlock(&self.cache)),
217        }
218    }
219}
220
221// --- Write Guards ---
222
223// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
224// as it is always possible to go back to an actual `RwLock` via `RwLockWriteGuard::rwlock`.
225// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
226// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
227pub struct StorageEngineWriteGuard<'a> {
228    store: parking_lot::RwLockWriteGuard<'a, ChunkStore>,
229    cache: parking_lot::RwLockWriteGuard<'a, QueryCache>,
230}
231
232impl<'a> StorageEngineWriteGuard<'a> {
233    #[inline]
234    pub fn downgrade(self) -> StorageEngineReadGuard<'a> {
235        StorageEngineReadGuard {
236            store: parking_lot::RwLockWriteGuard::downgrade(self.store),
237            cache: parking_lot::RwLockWriteGuard::downgrade(self.cache),
238        }
239    }
240}
241
242impl StorageEngineWriteGuard<'_> {
243    #[inline]
244    pub fn store(&mut self) -> &mut ChunkStore {
245        &mut self.store
246    }
247
248    #[inline]
249    pub fn cache(&mut self) -> &mut QueryCache {
250        &mut self.cache
251    }
252}
253
254// NOTE: None of these fields should ever be publicly exposed, either directly or through a method,
255// as it is always possible to go back to an actual `RwLock` via `ArcRwLockWriteGuard::rwlock`.
256// Doing so would defeat the deadlock protection that the `StorageEngine` offers.
257// Exposing references to the actual `ChunkStore` and `QueryCache` if ofc fine.
258pub struct StorageEngineArcWriteGuard {
259    store: parking_lot::ArcRwLockWriteGuard<parking_lot::RawRwLock, ChunkStore>,
260    cache: parking_lot::ArcRwLockWriteGuard<parking_lot::RawRwLock, QueryCache>,
261}
262
263impl StorageEngineArcWriteGuard {
264    #[inline]
265    pub fn downgrade(self) -> StorageEngineArcReadGuard {
266        StorageEngineArcReadGuard {
267            store: parking_lot::ArcRwLockWriteGuard::downgrade(self.store),
268            cache: parking_lot::ArcRwLockWriteGuard::downgrade(self.cache),
269        }
270    }
271}
272
273impl StorageEngineArcWriteGuard {
274    #[inline]
275    pub fn store(&mut self) -> &mut ChunkStore {
276        &mut self.store
277    }
278
279    #[inline]
280    pub fn cache(&mut self) -> &mut QueryCache {
281        &mut self.cache
282    }
283}
re_query/storage_engine.rs

re_query/
storage_engine.rs
